From d4dba7718a0550bf69f1462a14fee54b03490728 Mon Sep 17 00:00:00 2001
From: Julien Ruaux <jruaux@c64.arcillaux.local>
Date: Wed, 29 Jun 2022 18:43:09 -0700
Subject: [PATCH] fix: Support projecting non-indexed fields

---
 .../com/redis/trino/RediSearchSession.java    | 22 +++++++++++++++----
 .../TestRediSearchConnectorSmokeTest.java     |  2 +-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/subprojects/trino-redisearch/src/main/java/com/redis/trino/RediSearchSession.java b/subprojects/trino-redisearch/src/main/java/com/redis/trino/RediSearchSession.java
index d2047fb..4f9af92 100644
--- a/subprojects/trino-redisearch/src/main/java/com/redis/trino/RediSearchSession.java
+++ b/subprojects/trino-redisearch/src/main/java/com/redis/trino/RediSearchSession.java
@@ -12,6 +12,7 @@
 import static java.util.concurrent.TimeUnit.MINUTES;
 
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -29,6 +30,7 @@
 import com.redis.lettucemod.search.AggregateWithCursorResults;
 import com.redis.lettucemod.search.CreateOptions;
 import com.redis.lettucemod.search.CursorOptions;
+import com.redis.lettucemod.search.Document;
 import com.redis.lettucemod.search.Field;
 import com.redis.lettucemod.search.Group;
 import com.redis.lettucemod.search.IndexInfo;
@@ -163,13 +165,25 @@ private RediSearchTable loadTableSchema(SchemaTableName schemaTableName) {
 		if (indexInfo.isEmpty()) {
 			throw new TableNotFoundException(schemaTableName, format("Index '%s' not found", index), null);
 		}
+		Set<String> fields = new HashSet<>();
 		ImmutableList.Builder<RediSearchColumnHandle> columnHandles = ImmutableList.builder();
 		for (Field columnMetadata : indexInfo.get().getFields()) {
-			columnHandles.add(buildColumnHandle(columnMetadata));
+			RediSearchColumnHandle column = buildColumnHandle(columnMetadata);
+			fields.add(column.getName());
+			columnHandles.add(column);
 		}
-		RediSearchTableHandle tableHandle = new RediSearchTableHandle(RediSearchTableHandle.Type.SEARCH,
-				schemaTableName);
-		return new RediSearchTable(tableHandle, columnHandles.build());
+		SearchResults<String, String> results = connection.sync().search(index, "*");
+		for (Document<String, String> doc : results) {
+			for (String field : doc.keySet()) {
+				if (fields.contains(field)) {
+					continue;
+				}
+				columnHandles.add(new RediSearchColumnHandle(field, VarcharType.VARCHAR, false));
+				fields.add(field);
+			}
+		}
+		return new RediSearchTable(new RediSearchTableHandle(RediSearchTableHandle.Type.SEARCH, schemaTableName),
+				columnHandles.build());
 	}
 
 	private Optional<IndexInfo> indexInfo(String index) {
diff --git a/subprojects/trino-redisearch/src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java b/subprojects/trino-redisearch/src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java
index baa9877..a753c8e 100644
--- a/subprojects/trino-redisearch/src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java
+++ b/subprojects/trino-redisearch/src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java
@@ -106,7 +106,7 @@ protected QueryRunner createQueryRunner() throws Exception {
 	@Test
 	public void testNonIndexedFields() throws IOException {
 		Beers.populateIndex(redisearch.getTestContext().getConnection());
-		getQueryRunner().execute("select id, name from beers");
+		getQueryRunner().execute("select id, last_mod from beers");
 	}
 
 	@Test