From f06f3db647c81bc24fa9bdad33822ca6175c32eb Mon Sep 17 00:00:00 2001
From: shacharPash <93581407+shacharPash@users.noreply.github.com>
Date: Wed, 21 Dec 2022 15:11:33 +0200
Subject: [PATCH] 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
---
 redis/commands/search/query.py    | 8 ++++++++
 tests/test_asyncio/test_search.py | 9 +++++++++
 tests/test_search.py              | 9 +++++++++
 3 files changed, 26 insertions(+)

diff --git a/redis/commands/search/query.py b/redis/commands/search/query.py
index e51918f500..5071cfabf2 100644
--- a/redis/commands/search/query.py
+++ b/redis/commands/search/query.py
@@ -28,6 +28,7 @@ def __init__(self, query_string):
         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 @@ def slop(self, slop):
         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 @@ def _get_args_tags(self):
             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 88c80d301b..12313b65e3 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 08c76f7e08..12876f6262 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)