feat: Added support for DELETE statements

Julien Ruaux · Julien Ruaux · commit 4c1336bbc264 · 2022-12-30T14:27:21.000+08:00
diff --git a/pom.xml b/pom.xml
@@ -26,7 +26,7 @@
 		<github.owner>redis-field-engineering</github.owner>
 		<github.repo>redis-sql</github.repo>
 
-		<lettucemod.version>3.1.6</lettucemod.version>
+		<lettucemod.version>3.2.0</lettucemod.version>
 		<lettuce.version>6.2.2.RELEASE</lettuce.version>
 		<testcontainers-redis.version>1.6.2</testcontainers-redis.version>
 		<ulid.version>5.1.0</ulid.version>
diff --git a/src/main/java/com/redis/trino/RediSearchMetadata.java b/src/main/java/com/redis/trino/RediSearchMetadata.java
@@ -29,6 +29,8 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.airlift.slice.SliceUtf8.getCodePointAt;
 import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
+import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
+import static io.trino.spi.connector.RetryMode.NO_RETRIES;
 import static io.trino.spi.expression.StandardFunctions.LIKE_FUNCTION_NAME;
 import static java.util.Objects.requireNonNull;
 
@@ -244,6 +246,25 @@ public Optional<ConnectorOutputMetadata> finishInsert(ConnectorSession session,
 		return Optional.empty();
 	}
 
+	@Override
+	public RediSearchColumnHandle getDeleteRowIdColumnHandle(ConnectorSession session,
+			ConnectorTableHandle tableHandle) {
+		return RediSearchBuiltinField.ID.getColumnHandle();
+	}
+
+	@Override
+	public RediSearchTableHandle beginDelete(ConnectorSession session, ConnectorTableHandle tableHandle,
+			RetryMode retryMode) {
+		if (retryMode != NO_RETRIES) {
+			throw new TrinoException(NOT_SUPPORTED, "This connector does not support query retries");
+		}
+		return (RediSearchTableHandle) tableHandle;
+	}
+
+	@Override
+	public void finishDelete(ConnectorSession session, ConnectorTableHandle tableHandle, Collection<Slice> fragments) {
+	}
+
 	@Override
 	public ConnectorTableProperties getTableProperties(ConnectorSession session, ConnectorTableHandle table) {
 		RediSearchTableHandle handle = (RediSearchTableHandle) table;
diff --git a/src/main/java/com/redis/trino/RediSearchPageSink.java b/src/main/java/com/redis/trino/RediSearchPageSink.java
@@ -43,6 +43,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Random;
 import java.util.concurrent.CompletableFuture;
 
@@ -51,6 +52,9 @@
 import com.google.common.primitives.Shorts;
 import com.google.common.primitives.SignedBytes;
 import com.redis.lettucemod.api.StatefulRedisModulesConnection;
+import com.redis.lettucemod.search.CreateOptions;
+import com.redis.lettucemod.search.CreateOptions.DataType;
+import com.redis.lettucemod.search.IndexInfo;
 
 import io.airlift.slice.Slice;
 import io.lettuce.core.LettuceFutures;
@@ -78,26 +82,27 @@
 
 public class RediSearchPageSink implements ConnectorPageSink {
 
-	private final RediSearchSession rediSearchSession;
+	private final RediSearchSession session;
 	private final SchemaTableName schemaTableName;
 	private final List<RediSearchColumnHandle> columns;
 	private final UlidFactory factory = UlidFactory.newInstance(new Random());
 
 	public RediSearchPageSink(RediSearchSession rediSearchSession, SchemaTableName schemaTableName,
 			List<RediSearchColumnHandle> columns) {
-		this.rediSearchSession = rediSearchSession;
+		this.session = rediSearchSession;
 		this.schemaTableName = schemaTableName;
 		this.columns = columns;
 	}
 
 	@Override
 	public CompletableFuture<?> appendPage(Page page) {
-		StatefulRedisModulesConnection<String, String> connection = rediSearchSession.getConnection();
+		String prefix = prefix().orElse(schemaTableName.getTableName());
+		StatefulRedisModulesConnection<String, String> connection = session.getConnection();
 		connection.setAutoFlushCommands(false);
 		List<RedisFuture<?>> futures = new ArrayList<>();
 		for (int position = 0; position < page.getPositionCount(); position++) {
 			Map<String, String> map = new HashMap<>();
-			String key = schemaTableName.getTableName() + ":" + factory.create().toString();
+			String key = prefix + ":" + factory.create().toString();
 			for (int channel = 0; channel < page.getChannelCount(); channel++) {
 				RediSearchColumnHandle column = columns.get(channel);
 				Block block = page.getBlock(channel);
@@ -115,6 +120,31 @@ public CompletableFuture<?> appendPage(Page page) {
 		return NOT_BLOCKED;
 	}
 
+	private Optional<String> prefix() {
+		try {
+			RediSearchTable table = session.getTable(schemaTableName);
+			IndexInfo indexInfo = table.getIndexInfo();
+			CreateOptions<String, String> options = indexInfo.getIndexOptions();
+			Optional<DataType> on = options.getOn();
+			if (on.isEmpty() || on.get() != DataType.HASH) {
+				return Optional.empty();
+			}
+			if (options.getPrefixes().isEmpty()) {
+				return Optional.empty();
+			}
+			String prefix = options.getPrefixes().get(0);
+			if (prefix.equals("*")) {
+				return Optional.empty();
+			}
+			if (prefix.endsWith(":")) {
+				return Optional.of(prefix.substring(0, prefix.length() - 1));
+			}
+			return Optional.of(prefix);
+		} catch (Exception e) {
+			return Optional.empty();
+		}
+	}
+
 	private String getObjectValue(Type type, Block block, int position) {
 		if (type.equals(BooleanType.BOOLEAN)) {
 			return String.valueOf(type.getBoolean(block, position));
diff --git a/src/main/java/com/redis/trino/RediSearchPageSource.java b/src/main/java/com/redis/trino/RediSearchPageSource.java
@@ -27,40 +27,49 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.redis.lettucemod.search.Document;
 
+import io.airlift.slice.Slice;
 import io.airlift.slice.SliceOutput;
 import io.trino.spi.Page;
 import io.trino.spi.PageBuilder;
+import io.trino.spi.block.Block;
 import io.trino.spi.block.BlockBuilder;
-import io.trino.spi.connector.ConnectorPageSource;
+import io.trino.spi.connector.UpdatablePageSource;
 import io.trino.spi.type.Type;
 
-public class RediSearchPageSource implements ConnectorPageSource {
+public class RediSearchPageSource implements UpdatablePageSource {
 
 	private static final int ROWS_PER_REQUEST = 1024;
 
 	private final RediSearchPageSourceResultWriter writer = new RediSearchPageSourceResultWriter();
+	private final RediSearchSession session;
 	private final Iterator<Document<String, String>> cursor;
-	private final List<String> columnNames;
+	private final String[] columnNames;
 	private final List<Type> columnTypes;
+	private final PageBuilder pageBuilder;
+
 	private Document<String, String> currentDoc;
 	private long count;
 	private boolean finished;
 
-	private final PageBuilder pageBuilder;
-
-	public RediSearchPageSource(RediSearchSession rediSearchSession, RediSearchTableHandle tableHandle,
+	public RediSearchPageSource(RediSearchSession session, RediSearchTableHandle table,
 			List<RediSearchColumnHandle> columns) {
-		this.columnNames = columns.stream().map(RediSearchColumnHandle::getName).collect(Collectors.toUnmodifiableList());
-		this.columnTypes = columns.stream().map(RediSearchColumnHandle::getType).collect(Collectors.toUnmodifiableList());
-		this.cursor = rediSearchSession.search(tableHandle, columns).iterator();
+		this.session = session;
+		this.columnNames = columns.stream().map(RediSearchColumnHandle::getName).toArray(String[]::new);
+		this.columnTypes = columns.stream().map(RediSearchColumnHandle::getType)
+				.collect(Collectors.toUnmodifiableList());
+		this.cursor = session.search(table, columnNames).iterator();
 		this.currentDoc = null;
 		this.pageBuilder = new PageBuilder(columnTypes);
 	}
@@ -100,7 +109,7 @@ public Page getNextPage() {
 			pageBuilder.declarePosition();
 			for (int column = 0; column < columnTypes.size(); column++) {
 				BlockBuilder output = pageBuilder.getBlockBuilder(column);
-				String columnName = columnNames.get(column);
+				String columnName = columnNames[column];
 				String value = currentValue(columnName);
 				if (value == null) {
 					output.appendNull();
@@ -115,6 +124,17 @@ public Page getNextPage() {
 		return page;
 	}
 
+	@Override
+	public void deleteRows(Block rowIds) {
+		List<String> docIds = new ArrayList<>(rowIds.getPositionCount());
+		for (int i = 0; i < rowIds.getPositionCount(); i++) {
+			int len = rowIds.getSliceLength(i);
+			Slice slice = rowIds.getSlice(i, 0, len);
+			docIds.add(slice.toStringUtf8());
+		}
+		session.deleteDocs(docIds);
+	}
+
 	private String currentValue(String columnName) {
 		if (RediSearchBuiltinField.isBuiltinColumn(columnName)) {
 			if (RediSearchBuiltinField.ID.getName().equals(columnName)) {
@@ -127,6 +147,13 @@ private String currentValue(String columnName) {
 		return currentDoc.get(columnName);
 	}
 
+	@Override
+	public CompletableFuture<Collection<Slice>> finish() {
+		CompletableFuture<Collection<Slice>> future = new CompletableFuture<>();
+		future.complete(Collections.emptyList());
+		return future;
+	}
+
 	public static JsonGenerator createJsonGenerator(JsonFactory factory, SliceOutput output) throws IOException {
 		return factory.createGenerator((OutputStream) output);
 	}
diff --git a/src/main/java/com/redis/trino/RediSearchQueryBuilder.java b/src/main/java/com/redis/trino/RediSearchQueryBuilder.java
@@ -242,9 +242,9 @@ public Optional<Group> group(RediSearchTableHandle table) {
 		List<RediSearchAggregation> aggregates = table.getMetricAggregations();
 		List<String> groupFields = new ArrayList<>();
 		if (terms != null && !terms.isEmpty()) {
-			groupFields = terms.stream().map(RediSearchAggregationTerm::getTerm).collect(Collectors.toList());
+			groupFields = terms.stream().map(RediSearchAggregationTerm::getTerm).collect(Collectors.toUnmodifiableList());
 		}
-		List<Reducer> reducers = aggregates.stream().map(this::reducer).collect(Collectors.toList());
+		List<Reducer> reducers = aggregates.stream().map(this::reducer).collect(Collectors.toUnmodifiableList());
 		if (reducers.isEmpty()) {
 			return Optional.empty();
 		}
diff --git a/src/main/java/com/redis/trino/RediSearchSession.java b/src/main/java/com/redis/trino/RediSearchSession.java
@@ -34,8 +34,10 @@
 
 import java.io.File;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -48,10 +50,12 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.redis.lettucemod.api.StatefulRedisModulesConnection;
+import com.redis.lettucemod.search.AggregateOperation;
 import com.redis.lettucemod.search.AggregateWithCursorResults;
 import com.redis.lettucemod.search.CreateOptions;
 import com.redis.lettucemod.search.Document;
 import com.redis.lettucemod.search.Field;
+import com.redis.lettucemod.search.Group;
 import com.redis.lettucemod.search.IndexInfo;
 import com.redis.lettucemod.search.SearchResults;
 import com.redis.lettucemod.util.ClientBuilder;
@@ -161,6 +165,12 @@ public Set<String> getAllTables() throws SchemaNotFoundException {
 		return builder.build();
 	}
 
+	/**
+	 * 
+	 * @param schemaTableName SchemaTableName to load
+	 * @return RediSearchTable describing the RediSearch index
+	 * @throws TableNotFoundException if no index by that name was found
+	 */
 	public RediSearchTable getTable(SchemaTableName tableName) throws TableNotFoundException {
 		try {
 			return tableCache.getUnchecked(tableName);
@@ -172,13 +182,13 @@ public RediSearchTable getTable(SchemaTableName tableName) throws TableNotFoundE
 
 	@SuppressWarnings("unchecked")
 	public void createTable(SchemaTableName schemaTableName, List<RediSearchColumnHandle> columns) {
-		String tableName = schemaTableName.getTableName();
-		if (!connection.sync().ftList().contains(tableName)) {
+		String index = index(schemaTableName);
+		if (!connection.sync().ftList().contains(index)) {
 			List<Field<String>> fields = columns.stream().filter(c -> !c.getName().equals("_id"))
-					.map(c -> buildField(c.getName(), c.getType())).collect(Collectors.toList());
+					.map(c -> buildField(c.getName(), c.getType())).collect(Collectors.toUnmodifiableList());
 			CreateOptions.Builder<String, String> options = CreateOptions.<String, String>builder();
-			options.prefix(tableName + ":");
-			connection.sync().ftCreate(tableName, options.build(), fields.toArray(Field[]::new));
+			options.prefix(index + ":");
+			connection.sync().ftCreate(index, options.build(), fields.toArray(Field[]::new));
 		}
 	}
 
@@ -210,19 +220,26 @@ public void dropColumn(SchemaTableName schemaTableName, String columnName) {
 		throw new TrinoException(NOT_SUPPORTED, "This connector does not support dropping columns");
 	}
 
-	private RediSearchTable loadTableSchema(SchemaTableName schemaTableName) {
+	/**
+	 * 
+	 * @param schemaTableName SchemaTableName to load
+	 * @return RediSearchTable describing the RediSearch index
+	 * @throws TableNotFoundException if no index by that name was found
+	 */
+	private RediSearchTable loadTableSchema(SchemaTableName schemaTableName) throws TableNotFoundException {
 		String index = schemaTableName.getTableName();
-		Optional<IndexInfo> indexInfo = indexInfo(index);
-		if (indexInfo.isEmpty()) {
+		Optional<IndexInfo> indexInfoOptional = indexInfo(index);
+		if (indexInfoOptional.isEmpty()) {
 			throw new TableNotFoundException(schemaTableName, format("Index '%s' not found", index), null);
 		}
+		IndexInfo indexInfo = indexInfoOptional.get();
 		Set<String> fields = new HashSet<>();
 		ImmutableList.Builder<RediSearchColumnHandle> columns = ImmutableList.builder();
 		for (RediSearchBuiltinField builtinfield : RediSearchBuiltinField.values()) {
 			fields.add(builtinfield.getName());
 			columns.add(builtinfield.getColumnHandle());
 		}
-		for (Field<String> indexedField : indexInfo.get().getFields()) {
+		for (Field<String> indexedField : indexInfo.getFields()) {
 			RediSearchColumnHandle column = buildColumnHandle(indexedField);
 			fields.add(column.getName());
 			columns.add(column);
@@ -237,8 +254,9 @@ private RediSearchTable loadTableSchema(SchemaTableName schemaTableName) {
 				fields.add(docField);
 			}
 		}
-		return new RediSearchTable(new RediSearchTableHandle(RediSearchTableHandle.Type.SEARCH, schemaTableName),
-				columns.build());
+		RediSearchTableHandle tableHandle = new RediSearchTableHandle(RediSearchTableHandle.Type.SEARCH,
+				schemaTableName);
+		return new RediSearchTable(tableHandle, columns.build(), indexInfo);
 	}
 
 	private Optional<IndexInfo> indexInfo(String index) {
@@ -278,8 +296,7 @@ private Type columnType(TypeSignature typeSignature) {
 		return typeManager.fromSqlType(typeSignature.toString());
 	}
 
-	public SearchResults<String, String> search(RediSearchTableHandle tableHandle,
-			List<RediSearchColumnHandle> columns) {
+	public SearchResults<String, String> search(RediSearchTableHandle tableHandle, String[] columns) {
 		Search search = translator.search(tableHandle, columns);
 		log.info("Running %s", search);
 		return connection.sync().ftSearch(search.getIndex(), search.getQuery(), search.getOptions());
@@ -288,20 +305,32 @@ public SearchResults<String, String> search(RediSearchTableHandle tableHandle,
 	public AggregateWithCursorResults<String> aggregate(RediSearchTableHandle table) {
 		Aggregation aggregation = translator.aggregate(table);
 		log.info("Running %s", aggregation);
-		return connection.sync().ftAggregate(aggregation.getIndex(), aggregation.getQuery(),
-				aggregation.getCursorOptions(), aggregation.getOptions());
+		AggregateWithCursorResults<String> results = connection.sync().ftAggregate(aggregation.getIndex(),
+				aggregation.getQuery(), aggregation.getCursorOptions(), aggregation.getOptions());
+		List<AggregateOperation<?, ?>> groupBys = aggregation.getOptions().getOperations().stream()
+				.filter(o -> o.getType() == AggregateOperation.Type.GROUP).collect(Collectors.toUnmodifiableList());
+		if (results.isEmpty() && !groupBys.isEmpty()) {
+			Group groupBy = (Group) groupBys.get(0);
+			Optional<String> as = groupBy.getReducers()[0].getAs();
+			if (as.isPresent()) {
+				Map<String, Object> doc = new HashMap<>();
+				doc.put(as.get(), 0);
+				results.add(doc);
+			}
+		}
+		return results;
 	}
 
 	public AggregateWithCursorResults<String> cursorRead(RediSearchTableHandle tableHandle, long cursor) {
-		String index = index(tableHandle);
+		String index = index(tableHandle.getSchemaTableName());
 		if (config.getCursorCount() > 0) {
 			return connection.sync().ftCursorRead(index, cursor, config.getCursorCount());
 		}
 		return connection.sync().ftCursorRead(index, cursor);
 	}
 
-	private String index(RediSearchTableHandle tableHandle) {
-		return tableHandle.getSchemaTableName().getTableName();
+	private String index(SchemaTableName schemaTableName) {
+		return schemaTableName.getTableName();
 	}
 
 	private Field<String> buildField(String columnName, Type columnType) {
@@ -381,7 +410,11 @@ private TypeSignature varcharType() {
 	}
 
 	public void cursorDelete(RediSearchTableHandle tableHandle, long cursor) {
-		connection.sync().ftCursorDelete(index(tableHandle), cursor);
+		connection.sync().ftCursorDelete(index(tableHandle.getSchemaTableName()), cursor);
+	}
+
+	public Long deleteDocs(List<String> docIds) {
+		return connection.sync().del(docIds.toArray(String[]::new));
 	}
 
 }
diff --git a/src/main/java/com/redis/trino/RediSearchTable.java b/src/main/java/com/redis/trino/RediSearchTable.java
diff --git a/src/main/java/com/redis/trino/RediSearchTranslator.java b/src/main/java/com/redis/trino/RediSearchTranslator.java
diff --git a/src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java b/src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java

Original file line number	Diff line number	Diff line change
`@@ -242,9 +242,9 @@ public Optional<Group> group(RediSearchTableHandle table) {`
`242`	`242`	`List<RediSearchAggregation> aggregates = table.getMetricAggregations();`
`243`	`243`	`List<String> groupFields = new ArrayList<>();`
`244`	`244`	`if (terms != null && !terms.isEmpty()) {`
`245`		`- groupFields = terms.stream().map(RediSearchAggregationTerm::getTerm).collect(Collectors.toList());`
	`245`	`+ groupFields = terms.stream().map(RediSearchAggregationTerm::getTerm).collect(Collectors.toUnmodifiableList());`
`246`	`246`	`}`
`247`		`- List<Reducer> reducers = aggregates.stream().map(this::reducer).collect(Collectors.toList());`
	`247`	`+ List<Reducer> reducers = aggregates.stream().map(this::reducer).collect(Collectors.toUnmodifiableList());`
`248`	`248`	`if (reducers.isEmpty()) {`
`249`	`249`	`return Optional.empty();`
`250`	`250`	`}`