From 2cea637b2e936265d2001043d2b2d4e62559dc17 Mon Sep 17 00:00:00 2001
From: dvora-h <67596500+dvora-h@users.noreply.github.com>
Date: Thu, 4 Aug 2022 16:46:28 +0300
Subject: [PATCH] Add support for WITHSUFFIXTRIE to FT.CREATE (#2324)

* withsuffixtrie

* Update test_search.py

* fix
---
 redis/commands/search/field.py    | 12 +++++++++++-
 tests/test_asyncio/test_search.py | 24 ++++++++++++++++++++++++
 tests/test_search.py              | 24 ++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/redis/commands/search/field.py b/redis/commands/search/field.py
index 89ed97357a..6f31ce1fc2 100644
--- a/redis/commands/search/field.py
+++ b/redis/commands/search/field.py
@@ -64,6 +64,7 @@ def __init__(
         weight: float = 1.0,
         no_stem: bool = False,
         phonetic_matcher: str = None,
+        withsuffixtrie: bool = False,
         **kwargs,
     ):
         Field.__init__(self, name, args=[Field.TEXT, Field.WEIGHT, weight], **kwargs)
@@ -78,6 +79,8 @@ def __init__(
         ]:
             Field.append_arg(self, self.PHONETIC)
             Field.append_arg(self, phonetic_matcher)
+        if withsuffixtrie:
+            Field.append_arg(self, "WITHSUFFIXTRIE")
 
 
 class NumericField(Field):
@@ -108,11 +111,18 @@ class TagField(Field):
     CASESENSITIVE = "CASESENSITIVE"
 
     def __init__(
-        self, name: str, separator: str = ",", case_sensitive: bool = False, **kwargs
+        self,
+        name: str,
+        separator: str = ",",
+        case_sensitive: bool = False,
+        withsuffixtrie: bool = False,
+        **kwargs,
     ):
         args = [Field.TAG, self.SEPARATOR, separator]
         if case_sensitive:
             args.append(self.CASESENSITIVE)
+        if withsuffixtrie:
+            args.append("WITHSUFFIXTRIE")
 
         Field.__init__(self, name, args=args, **kwargs)
 
diff --git a/tests/test_asyncio/test_search.py b/tests/test_asyncio/test_search.py
index 4824e49624..862c21b1c3 100644
--- a/tests/test_asyncio/test_search.py
+++ b/tests/test_asyncio/test_search.py
@@ -1045,6 +1045,30 @@ async def test_aggregations_sort_by_and_limit(modclient: redis.Redis):
     assert res.rows[0] == ["t1", "b"]
 
 
+@pytest.mark.redismod
+@pytest.mark.experimental
+async def test_withsuffixtrie(modclient: redis.Redis):
+    # create index
+    assert await modclient.ft().create_index((TextField("txt"),))
+    await waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
+    info = await modclient.ft().info()
+    assert "WITHSUFFIXTRIE" not in info["attributes"][0]
+    assert await modclient.ft().dropindex("idx")
+
+    # create withsuffixtrie index (text field)
+    assert await modclient.ft().create_index((TextField("t", withsuffixtrie=True)))
+    await waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
+    info = await modclient.ft().info()
+    assert "WITHSUFFIXTRIE" in info["attributes"][0]
+    assert await modclient.ft().dropindex("idx")
+
+    # create withsuffixtrie index (tag field)
+    assert await modclient.ft().create_index((TagField("t", withsuffixtrie=True)))
+    await waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
+    info = await modclient.ft().info()
+    assert "WITHSUFFIXTRIE" in info["attributes"][0]
+
+
 @pytest.mark.redismod
 @skip_if_redis_enterprise()
 async def test_search_commands_in_pipeline(modclient: redis.Redis):
diff --git a/tests/test_search.py b/tests/test_search.py
index ee4fa55f68..5fe5ab17fc 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -1712,3 +1712,27 @@ def test_expire_while_search(modclient: redis.Redis):
         modclient.ft().search(Query("*")).docs[1]
     time.sleep(1)
     assert 2 == modclient.ft().search(Query("*")).total
+
+
+@pytest.mark.redismod
+@pytest.mark.experimental
+def test_withsuffixtrie(modclient: redis.Redis):
+    # create index
+    assert modclient.ft().create_index((TextField("txt"),))
+    waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
+    info = modclient.ft().info()
+    assert "WITHSUFFIXTRIE" not in info["attributes"][0]
+    assert modclient.ft().dropindex("idx")
+
+    # create withsuffixtrie index (text fiels)
+    assert modclient.ft().create_index((TextField("t", withsuffixtrie=True)))
+    waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
+    info = modclient.ft().info()
+    assert "WITHSUFFIXTRIE" in info["attributes"][0]
+    assert modclient.ft().dropindex("idx")
+
+    # create withsuffixtrie index (tag field)
+    assert modclient.ft().create_index((TagField("t", withsuffixtrie=True)))
+    waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
+    info = modclient.ft().info()
+    assert "WITHSUFFIXTRIE" in info["attributes"][0]