Add TIMEOUT to query class (#2519)

* add timeout to query class

* Add test_timeout

* fix lines

* fix format

* add test & fixes

* merge tests

* change timeout to not_a_number

* change q1 to q2

* Fix async method

3 files changed, 26 insertions, 0 deletions

diff --git a/redis/commands/search/query.py b/redis/commands/search/query.py
index e51918f..5071cfa 100644
--- a/redis/commands/search/query.py
+++ b/redis/commands/search/query.py
@@ -28,6 +28,7 @@ class Query:
         self._filters = list()
         self._ids = None
         self._slop = -1
+        self._timeout = None
         self._in_order = False
         self._sortby = None
         self._return_fields = []
@@ -131,6 +132,11 @@ class Query:
         self._slop = slop
         return self
 
+    def timeout(self, timeout):
+        """overrides the timeout parameter of the module"""
+        self._timeout = timeout
+        return self
+
     def in_order(self):
         """
         Match only documents where the query terms appear in
@@ -188,6 +194,8 @@ class Query:
             args += self._ids
         if self._slop >= 0:
             args += ["SLOP", self._slop]
+        if self._timeout:
+            args += ["TIMEOUT", self._timeout]
         if self._in_order:
             args.append("INORDER")
         if self._return_fields:
diff --git a/tests/test_asyncio/test_search.py b/tests/test_asyncio/test_search.py
index 88c80d3..12313b6 100644
--- a/tests/test_asyncio/test_search.py
+++ b/tests/test_asyncio/test_search.py
@@ -1001,3 +1001,12 @@ async def test_search_commands_in_pipeline(modclient: redis.Redis):
     assert "doc2" == res[3][4]
     assert res[3][5] is None
     assert res[3][3] == res[3][6] == ["txt", "foo bar"]
+
+
+@pytest.mark.redismod
+async def test_query_timeout(modclient: redis.Redis):
+    q1 = Query("foo").timeout(5000)
+    assert q1.get_args() == ["foo", "TIMEOUT", 5000, "LIMIT", 0, 10]
+    q2 = Query("foo").timeout("not_a_number")
+    with pytest.raises(redis.ResponseError):
+        await modclient.ft().search(q2)
diff --git a/tests/test_search.py b/tests/test_search.py
index 08c76f7..12876f6 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -1615,3 +1615,12 @@ def test_withsuffixtrie(modclient: redis.Redis):
     waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
     info = modclient.ft().info()
     assert "WITHSUFFIXTRIE" in info["attributes"][0]
+
+
+@pytest.mark.redismod
+def test_query_timeout(modclient: redis.Redis):
+    q1 = Query("foo").timeout(5000)
+    assert q1.get_args() == ["foo", "TIMEOUT", 5000, "LIMIT", 0, 10]
+    q2 = Query("foo").timeout("not_a_number")
+    with pytest.raises(redis.ResponseError):
+        modclient.ft().search(q2)