From 6dc5c645e4388ad78f483d0f7ab6ec38b8b125f0 Mon Sep 17 00:00:00 2001
From: Julien Ruaux <jrx@github.com>
Date: Thu, 5 Jan 2023 15:20:01 +0800
Subject: [PATCH] feat: Added config option to specific protocol version.
 Resolves #7

---
 src/docs/asciidoc/index.adoc                         |  4 ++++
 src/main/java/com/redis/trino/RediSearchConfig.java  | 12 ++++++++++++
 src/main/java/com/redis/trino/RediSearchSession.java |  4 ++++
 .../java/com/redis/trino/TestRediSearchConfig.java   | 11 ++++++-----
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/docs/asciidoc/index.adoc b/src/docs/asciidoc/index.adoc
index c88d767..6288990 100644
--- a/src/docs/asciidoc/index.adoc
+++ b/src/docs/asciidoc/index.adoc
@@ -82,6 +82,10 @@ redisearch.uri=redis://localhost:6379
 |Connect to a Redis Cluster.
 |`false`
 
+|`redisearch.resp2`
+|Force Redis protocol version to RESP2.
+|`false`
+
 |===
 
 The RediSearch connector provides additional security options to support Redis servers with TLS mode.
diff --git a/src/main/java/com/redis/trino/RediSearchConfig.java b/src/main/java/com/redis/trino/RediSearchConfig.java
index a9b723a..2122532 100644
--- a/src/main/java/com/redis/trino/RediSearchConfig.java
+++ b/src/main/java/com/redis/trino/RediSearchConfig.java
@@ -57,6 +57,7 @@ public class RediSearchConfig {
 	private String keyPath;
 	private String certPath;
 	private String keyPassword;
+	private boolean resp2;
 
 	@Min(0)
 	public long getCursorCount() {
@@ -91,6 +92,17 @@ public RediSearchConfig setCaseInsensitiveNames(boolean caseInsensitiveNames) {
 		return this;
 	}
 
+	public boolean isResp2() {
+		return resp2;
+	}
+
+	@Config("redisearch.resp2")
+	@ConfigDescription("Force Redis protocol version to RESP2")
+	public RediSearchConfig setResp2(boolean resp2) {
+		this.resp2 = resp2;
+		return this;
+	}
+
 	@Config("redisearch.table-cache-expiration")
 	@ConfigDescription("Duration in seconds since the entry creation after which a table should be automatically removed from the cache.")
 	public RediSearchConfig setTableCacheExpiration(long expirationDuration) {
diff --git a/src/main/java/com/redis/trino/RediSearchSession.java b/src/main/java/com/redis/trino/RediSearchSession.java
index 5d2a8c6..8445f5a 100644
--- a/src/main/java/com/redis/trino/RediSearchSession.java
+++ b/src/main/java/com/redis/trino/RediSearchSession.java
@@ -68,6 +68,7 @@
 import io.lettuce.core.AbstractRedisClient;
 import io.lettuce.core.RedisURI;
 import io.lettuce.core.SslVerifyMode;
+import io.lettuce.core.protocol.ProtocolVersion;
 import io.trino.spi.HostAddress;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ColumnMetadata;
@@ -122,6 +123,9 @@ private AbstractRedisClient client(RediSearchConfig config) {
 		config.getCertPath().map(File::new).ifPresent(builder::keyCert);
 		config.getKeyPassword().ifPresent(p -> builder.keyPassword(p.toCharArray()));
 		config.getCaCertPath().map(File::new).ifPresent(builder::trustManager);
+		if (config.isResp2()) {
+			builder.protocolVersion(ProtocolVersion.RESP2);
+		}
 		return builder.build();
 	}
 
diff --git a/src/test/java/com/redis/trino/TestRediSearchConfig.java b/src/test/java/com/redis/trino/TestRediSearchConfig.java
index 386b444..250001a 100644
--- a/src/test/java/com/redis/trino/TestRediSearchConfig.java
+++ b/src/test/java/com/redis/trino/TestRediSearchConfig.java
@@ -2,10 +2,10 @@
 
 import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
 import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
-import static org.testng.Assert.assertEquals;
 
 import java.util.Map;
 
+import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import com.google.common.collect.ImmutableMap;
@@ -17,7 +17,7 @@ public class TestRediSearchConfig {
 	@Test
 	public void testDefaults() {
 		assertRecordedDefaults(recordDefaults(RediSearchConfig.class).setUri(null).setInsecure(false).setUsername(null)
-				.setPassword(null).setDefaultSchema(RediSearchConfig.DEFAULT_SCHEMA)
+				.setResp2(false).setPassword(null).setDefaultSchema(RediSearchConfig.DEFAULT_SCHEMA)
 				.setDefaultLimit(RediSearchConfig.DEFAULT_LIMIT).setCaseInsensitiveNames(false)
 				.setCursorCount(RediSearchConfig.DEFAULT_CURSOR_COUNT)
 				.setTableCacheExpiration(RediSearchConfig.DEFAULT_TABLE_CACHE_EXPIRATION.toSeconds())
@@ -30,15 +30,16 @@ public void testExplicitPropertyMappings() {
 		String uri = "redis://redis.example.com:12000";
 		String defaultSchema = "myschema";
 		Map<String, String> properties = new ImmutableMap.Builder<String, String>().put("redisearch.uri", uri)
-				.put("redisearch.default-schema-name", defaultSchema).build();
+				.put("redisearch.default-schema-name", defaultSchema).put("redisearch.resp2", "true").build();
 
 		ConfigurationFactory configurationFactory = new ConfigurationFactory(properties);
 		RediSearchConfig config = configurationFactory.build(RediSearchConfig.class);
 
 		RediSearchConfig expected = new RediSearchConfig().setDefaultSchema(defaultSchema).setUri(uri);
 
-		assertEquals(config.getDefaultSchema(), expected.getDefaultSchema());
-		assertEquals(config.getUri(), expected.getUri());
+		Assert.assertEquals(config.getDefaultSchema(), expected.getDefaultSchema());
+		Assert.assertEquals(config.getUri(), expected.getUri());
+		Assert.assertTrue(config.isResp2());
 	}
 
 }