Support RediSearch Suggestion commands (#3031)

The commands are: - FT.SUGADD - FT.SUGGET - FT.SUGDEL - FT.SUGLEN
redis · Jun 19, 2022 · 645f199 · 645f199
1 parent cbbfe11
commit 645f199
Show file tree

Hide file tree

Showing 6 changed files with 214 additions and 189 deletions.
diff --git a/src/main/java/redis/clients/jedis/CommandObjects.java b/src/main/java/redis/clients/jedis/CommandObjects.java
@@ -3142,6 +3142,45 @@ public final CommandObject<String> ftConfigSet(String option, String value) {
   public CommandObject<String> ftConfigSet(String indexName, String option, String value) {
     return ftConfigSet(option, value);
   }
+
+  public final CommandObject<Long> ftSugAdd(String key, String string, double score) {
+    return new CommandObject<>(commandArguments(SearchCommand.SUGADD).key(key).add(string).add(score), BuilderFactory.LONG);
+  }
+
+  public final CommandObject<Long> ftSugAddIncr(String key, String string, double score) {
+    return new CommandObject<>(commandArguments(SearchCommand.SUGADD).key(key).add(string).add(score).add(SearchKeyword.INCR), BuilderFactory.LONG);
+  }
+
+  public final CommandObject<List<String>> ftSugGet(String key, String prefix) {
+    return new CommandObject<>(commandArguments(SearchCommand.SUGGET).key(key).add(prefix), BuilderFactory.STRING_LIST);
+  }
+
+  public final CommandObject<List<String>> ftSugGet(String key, String prefix, boolean fuzzy, int max) {
+    CommandArguments args = commandArguments(SearchCommand.SUGGET).key(key).add(prefix);
+    if (fuzzy) args.add(SearchKeyword.FUZZY);
+    args.add(SearchKeyword.MAX).add(max);
+    return new CommandObject<>(args, BuilderFactory.STRING_LIST);
+  }
+
+  public final CommandObject<List<Tuple>> ftSugGetWithScores(String key, String prefix) {
+    return new CommandObject<>(commandArguments(SearchCommand.SUGGET).key(key).add(prefix).add(SearchKeyword.WITHSCORES), BuilderFactory.TUPLE_LIST);
+  }
+
+  public final CommandObject<List<Tuple>> ftSugGetWithScores(String key, String prefix, boolean fuzzy, int max) {
+    CommandArguments args = commandArguments(SearchCommand.SUGGET).key(key).add(prefix);
+    if (fuzzy) args.add(SearchKeyword.FUZZY);
+    args.add(SearchKeyword.MAX).add(max);
+    args.add(SearchKeyword.WITHSCORES);
+    return new CommandObject<>(args, BuilderFactory.TUPLE_LIST);
+  }
+
+  public final CommandObject<Boolean> ftSugDel(String key, String string) {
+    return new CommandObject<>(commandArguments(SearchCommand.SUGDEL).key(key).add(string), BuilderFactory.BOOLEAN);
+  }
+
+  public final CommandObject<Long> ftSugLen(String key) {
+    return new CommandObject<>(commandArguments(SearchCommand.SUGLEN).key(key), BuilderFactory.LONG);
+  }
   // RediSearch commands
 
   // RedisJSON commands

diff --git a/src/main/java/redis/clients/jedis/UnifiedJedis.java b/src/main/java/redis/clients/jedis/UnifiedJedis.java
@@ -3521,6 +3521,46 @@ public String ftConfigSet(String option, String value) {
   public String ftConfigSet(String indexName, String option, String value) {
     return executeCommand(commandObjects.ftConfigSet(indexName, option, value));
   }
+
+  @Override
+  public long ftSugAdd(String key, String string, double score) {
+    return executeCommand(commandObjects.ftSugAdd(key, string, score));
+  }
+
+  @Override
+  public long ftSugAddIncr(String key, String string, double score) {
+    return executeCommand(commandObjects.ftSugAddIncr(key, string, score));
+  }
+
+  @Override
+  public List<String> ftSugGet(String key, String prefix) {
+    return executeCommand(commandObjects.ftSugGet(key, prefix));
+  }
+
+  @Override
+  public List<String> ftSugGet(String key, String prefix, boolean fuzzy, int max) {
+    return executeCommand(commandObjects.ftSugGet(key, prefix, fuzzy, max));
+  }
+
+  @Override
+  public List<Tuple> ftSugGetWithScores(String key, String prefix) {
+    return executeCommand(commandObjects.ftSugGetWithScores(key, prefix));
+  }
+
+  @Override
+  public List<Tuple> ftSugGetWithScores(String key, String prefix, boolean fuzzy, int max) {
+    return executeCommand(commandObjects.ftSugGetWithScores(key, prefix, fuzzy, max));
+  }
+
+  @Override
+  public boolean ftSugDel(String key, String string) {
+    return executeCommand(commandObjects.ftSugDel(key, string));
+  }
+
+  @Override
+  public long ftSugLen(String key) {
+    return executeCommand(commandObjects.ftSugLen(key));
+  }
   // RediSearch commands
 
   // RedisJSON commands

diff --git a/src/main/java/redis/clients/jedis/search/RediSearchCommands.java b/src/main/java/redis/clients/jedis/search/RediSearchCommands.java
@@ -2,6 +2,7 @@
 
 import java.util.List;
 import java.util.Map;
+import redis.clients.jedis.resps.Tuple;
 import redis.clients.jedis.search.aggr.AggregationBuilder;
 import redis.clients.jedis.search.aggr.AggregationResult;
 
@@ -52,4 +53,20 @@ default String ftAlter(String indexName, Schema.Field... fields) {
   String ftConfigSet(String option, String value);
 
   String ftConfigSet(String indexName, String option, String value);
+
+  long ftSugAdd(String key, String string, double score);
+
+  long ftSugAddIncr(String key, String string, double score);
+
+  List<String> ftSugGet(String key, String prefix);
+
+  List<String> ftSugGet(String key, String prefix, boolean fuzzy, int max);
+
+  List<Tuple> ftSugGetWithScores(String key, String prefix);
+
+  List<Tuple> ftSugGetWithScores(String key, String prefix, boolean fuzzy, int max);
+
+  boolean ftSugDel(String key, String string);
+
+  long ftSugLen(String key);
 }
diff --git a/src/main/java/redis/clients/jedis/search/SearchProtocol.java b/src/main/java/redis/clients/jedis/search/SearchProtocol.java
@@ -22,10 +22,10 @@ public enum SearchCommand implements ProtocolCommand {
     ALIASDEL("FT.ALIASDEL"),
     SYNUPDATE("FT.SYNUPDATE"),
     SYNDUMP("FT.SYNDUMP"),
-//    SUGADD("FT.SUGADD"),
-//    SUGGET("FT.SUGGET"),
-//    SUGDEL("FT.SUGDEL"),
-//    SUGLEN("FT.SUGLEN"),
+    SUGADD("FT.SUGADD"),
+    SUGGET("FT.SUGGET"),
+    SUGDEL("FT.SUGDEL"),
+    SUGLEN("FT.SUGLEN"),
     DROPINDEX("FT.DROPINDEX");
 
     private final byte[] raw;

diff --git a/src/test/java/redis/clients/jedis/modules/search/SearchTest.java b/src/test/java/redis/clients/jedis/modules/search/SearchTest.java
@@ -1057,191 +1057,6 @@ public void testHighlightSummarize() throws Exception {
         res.getDocuments().get(0).get("text"));
 
   }
-//
-//  @Test
-//  public void testLanguage() throws Exception {
-//    Schema sc = new Schema().addTextField("text", 1.0);
-//    client.ftCreate(index, IndexOptions.defaultOptions(), sc);
-//
-//    Document d = new Document("doc1").set("text", "hello");
-//    AddOptions options = new AddOptions().setLanguage("spanish");
-//    assertTrue(client.addDocument(d, options));
-//    boolean caught = false;
-//
-//    options.setLanguage("ybreski");
-//    client.deleteDocument(d.getId());
-//
-//    try {
-//      client.addDocument(d, options);
-//    } catch (JedisDataException t) {
-//      caught = true;
-//    }
-//    assertTrue(caught);
-//  }
-//
-//  @Test
-//  public void testGet() throws Exception {
-//    client.ftCreate(index, IndexOptions.defaultOptions(), new Schema().addTextField("txt1", 1.0));
-//    client.addDocument(new Document("doc1").set("txt1", "Hello World!"), new AddOptions());
-//    Document d = client.getDocument("doc1");
-//    assertNotNull(d);
-//    assertEquals("Hello World!", d.get("txt1"));
-//
-//    // Get something that does not exist. Shouldn't explode
-//    assertNull(client.getDocument("nonexist"));
-//
-//    // Test decode=false mode
-//    d = client.getDocument("doc1", false);
-//    assertNotNull(d);
-//    assertTrue(Arrays.equals(SafeEncoder.encode("Hello World!"), (byte[]) d.get("txt1")));
-//  }
-//
-//  @Test
-//  public void testMGet() throws Exception {
-//    client.ftCreate(index, IndexOptions.defaultOptions(), new Schema().addTextField("txt1", 1.0));
-//    client.addDocument(new Document("doc1").set("txt1", "Hello World!1"), new AddOptions());
-//    client.addDocument(new Document("doc2").set("txt1", "Hello World!2"), new AddOptions());
-//    client.addDocument(new Document("doc3").set("txt1", "Hello World!3"), new AddOptions());
-//
-//    List<Document> docs = client.getDocuments();
-//    assertEquals(0, docs.size());
-//
-//    docs = client.getDocuments("doc1", "doc3", "doc4");
-//    assertEquals(3, docs.size());
-//    assertEquals("Hello World!1", docs.get(0).get("txt1"));
-//    assertEquals("Hello World!3", docs.get(1).get("txt1"));
-//    assertNull(docs.get(2));
-//
-//    // Test decode=false mode
-//    docs = client.getDocuments(false, "doc2");
-//    assertEquals(1, docs.size());
-//    assertTrue(Arrays.equals(SafeEncoder.encode("Hello World!2"), (byte[]) docs.get(0).get("txt1")));
-//  }
-//
-//  private static void assertUnknownIndex(JedisDataException jde) {
-//    assertTrue(jde.getMessage().toLowerCase().contains("unknown index"));
-//  }
-//
-//  @Test
-//  public void testAddSuggestionGetSuggestionFuzzy() throws Exception {
-//    Suggestion suggestion = Suggestion.builder().str("TOPIC OF WORDS").score(1).build();
-//    // test can add a suggestion string
-//    assertTrue(suggestion.toString() + " insert should of returned at least 1", client.addSuggestion(suggestion, true) > 0);
-//    // test that the partial part of that string will be returned using fuzzy
-//
-//    assertEquals(suggestion.toString() + " suppose to be returned", suggestion, client.getSuggestion(suggestion.getString().substring(0, 3), SuggestionOptions.builder().build()).get(0));
-//  }
-//
-//  @Test
-//  public void testAddSuggestionGetSuggestion() throws Exception {
-//    try {
-//      Suggestion.builder().str("ANOTHER_WORD").score(3).build();
-//      fail("Illegal score");
-//    } catch (IllegalStateException e) {
-//    }
-//
-//    try {
-//      Suggestion.builder().score(1).build();
-//      fail("Missing required string");
-//    } catch (IllegalStateException e) {
-//    }
-//
-//    Suggestion suggestion = Suggestion.builder().str("ANOTHER_WORD").score(1).build();
-//    Suggestion noMatch = Suggestion.builder().str("_WORD MISSED").score(1).build();
-//
-//    assertTrue(suggestion.toString() + " should of inserted at least 1", client.addSuggestion(suggestion, false) > 0);
-//    assertTrue(noMatch.toString() + " should of inserted at least 1", client.addSuggestion(noMatch, false) > 0);
-//
-//    // test that with a partial part of that string will have the entire word returned SuggestionOptions.builder().build()
-//    assertEquals(suggestion.getString() + " did not get a match with 3 characters", 1, client.getSuggestion(suggestion.getString().substring(0, 3), SuggestionOptions.builder().fuzzy().build()).size());
-//
-//    // turn off fuzzy start at second word no hit
-//    assertEquals(noMatch.getString() + " no fuzzy and starting at 1, should not match", 0, client.getSuggestion(noMatch.getString().substring(1, 6), SuggestionOptions.builder().build()).size());
-//    // my attempt to trigger the fuzzy by 1 character
-//    assertEquals(noMatch.getString() + " fuzzy is on starting at 1 position should match", 1, client.getSuggestion(noMatch.getString().substring(1, 6), SuggestionOptions.builder().fuzzy().build()).size());
-//  }
-//
-//  @Test
-//  public void testAddSuggestionGetSuggestionPayloadScores() throws Exception {
-//    Suggestion suggestion = Suggestion.builder().str("COUNT_ME TOO").payload("PAYLOADS ROCK ").score(0.2).build();
-//    assertTrue(suggestion.toString() + " insert should of at least returned 1", client.addSuggestion(suggestion, false) > 0);
-//    assertTrue("Count single added should return more than 1", client.addSuggestion(suggestion.toBuilder().str("COUNT").payload("My PAYLOAD is better").build(), false) > 1);
-//    assertTrue("Count single added should return more than 1", client.addSuggestion(suggestion.toBuilder().str("COUNT_ANOTHER").score(1).payload(null).build(), false) > 1);
-//
-//    Suggestion noScoreOrPayload = Suggestion.builder().str("COUNT NO PAYLOAD OR COUNT").build();
-//    assertTrue("Count single added should return more than 1", client.addSuggestion(noScoreOrPayload, true) > 1);
-//
-//    List<Suggestion> payloads = client.getSuggestion(suggestion.getString().substring(0, 3), SuggestionOptions.builder().with(SuggestionOptions.With.PAYLOAD_AND_SCORES).build());
-//    assertEquals("4 suggestions with scores and payloads ", 4, payloads.size());
-//    assertTrue("Assert that a suggestion has a payload ", payloads.get(2).getPayload().length() > 0);
-//    assertTrue("Assert that a suggestion has a score not default 1 ", payloads.get(1).getScore() < .299);
-//
-//  }
-//
-//  @Test
-//  public void testAddSuggestionGetSuggestionPayload() throws Exception {
-//    client.addSuggestion(Suggestion.builder().str("COUNT_ME TOO").payload("PAYLOADS ROCK ").build(), false);
-//    client.addSuggestion(Suggestion.builder().str("COUNT").payload("ANOTHER PAYLOAD ").build(), false);
-//    client.addSuggestion(Suggestion.builder().str("COUNTNO PAYLOAD OR COUNT").build(), false);
-//
-//    // test that with a partial part of that string will have the entire word returned
-//    List<Suggestion> payloads = client.getSuggestion("COU", SuggestionOptions.builder().max(3).fuzzy().with(SuggestionOptions.With.PAYLOAD).build());
-//    assertEquals("3 suggestions payloads ", 3, payloads.size());
-//
-//  }
-//
-//  @Test
-//  public void testGetSuggestionNoPayloadTwoOnly() throws Exception {
-//    client.addSuggestion(Suggestion.builder().str("DIFF_WORD").score(0.4).payload("PAYLOADS ROCK ").build(), false);
-//    client.addSuggestion(Suggestion.builder().str("DIFF wording").score(0.5).payload("ANOTHER PAYLOAD ").build(), false);
-//    client.addSuggestion(Suggestion.builder().str("DIFFERENT").score(0.7).payload("I am a payload").build(), false);
-//
-//    List<Suggestion> payloads = client.getSuggestion("DIF", SuggestionOptions.builder().max(2).build());
-//    assertEquals("3 suggestions should match but only asking for 2 and payloads should have 2 items in array", 2, payloads.size());
-//
-//    List<Suggestion> three = client.getSuggestion("DIF", SuggestionOptions.builder().max(3).build());
-//    assertEquals("3 suggestions and payloads should have 3 items in array", 3, three.size());
-//
-//  }
-//
-//  @Test
-//  public void testGetSuggestionWithScore() throws Exception {
-//    client.addSuggestion(Suggestion.builder().str("DIFF_WORD").score(0.4).payload("PAYLOADS ROCK ").build(), true);
-//    List<Suggestion> list = client.getSuggestion("DIF", SuggestionOptions.builder().max(2).with(SuggestionOptions.With.SCORES).build());
-//    assertTrue(list.get(0).getScore() <= .2);
-//  }
-//
-//  @Test
-//  public void testGetSuggestionAllNoHit() throws Exception {
-//    client.addSuggestion(Suggestion.builder().str("NO WORD").score(0.4).build(), false);
-//
-//    List<Suggestion> none = client.getSuggestion("DIF", SuggestionOptions.builder().max(3).with(SuggestionOptions.With.SCORES).build());
-//    assertEquals("Empty list not hit in index for partial word", 0, none.size());
-//  }
-//
-//  @Test
-//  public void testAddSuggestionDeleteSuggestionLength() throws Exception {
-//    client.addSuggestion(Suggestion.builder().str("TOPIC OF WORDS").score(1).build(), true);
-//    client.addSuggestion(Suggestion.builder().str("ANOTHER ENTRY").score(1).build(), true);
-//
-//    long result = client.deleteSuggestion("ANOTHER ENTRY");
-//    assertEquals("The delete of the suggestion should return 1", 1, result);
-//    assertEquals(1L, client.getSuggestionLength().longValue());
-//
-//    result = client.deleteSuggestion("ANOTHER ENTRY THAT IS NOT PRESENT");
-//    assertEquals("The delete of the suggestion should return 0", 0, result);
-//    assertEquals(1L, client.getSuggestionLength().longValue());
-//  }
-//
-//  @Test
-//  public void testAddSuggestionGetSuggestionLength() throws Exception {
-//    client.addSuggestion(Suggestion.builder().str("TOPIC OF WORDS").score(1).build(), true);
-//    client.addSuggestion(Suggestion.builder().str("ANOTHER ENTRY").score(1).build(), true);
-//    assertEquals(2L, client.getSuggestionLength().longValue());
-//
-//    client.addSuggestion(Suggestion.builder().str("FINAL ENTRY").score(1).build(), true);
-//    assertEquals(3L, client.getSuggestionLength().longValue());
-//  }
 
   @Test
   public void getTagField() {