Add optional timeout argument to graph queries (#87)

* Add optional timeout argument to graph queries

* Address PR comments

Author: jeffreylovitz

README.md

@@ -51,6 +51,8 @@ result = redis_graph.query(query, params)
 # Print resultset
 result.pretty_print()
 
+# Use query timeout to raise an exception if the query takes over 10 milliseconds
+result = redis_graph.query(query, params, timeout=10)
 
 # Iterate through resultset
 for record in result.result_set:

redisgraph/graph.py

@@ -117,14 +117,19 @@ def build_params_header(self, params):
             params_header += str(key) + "=" + str(value) + " "
         return params_header
 
-    def query(self, q, params=None):
+    def query(self, q, params=None, timeout=None):
         """
         Executes a query against the graph.
         """
         if params is not None:
             q = self.build_params_header(params) + q
 
-        response = self.redis_con.execute_command("GRAPH.QUERY", self.name, q, "--compact")
+        command = ["GRAPH.QUERY", self.name, q, "--compact"]
+        if timeout:
+            if not isinstance(timeout, int):
+                raise Exception("Timeout argument must be a positive integer")
+            command += ["timeout", timeout]
+        response = self.redis_con.execute_command(*command)
         return QueryResult(self, response)
 
     def _execution_plan_to_string(self, plan):

test.py

@@ -218,5 +218,26 @@ def test_execution_plan(self):
 
         redis_graph.delete()
 
+    def test_query_timeout(self):
+        redis_graph = Graph('timeout', self.r)
+        # Build a sample graph with 1000 nodes.
+        redis_graph.query("UNWIND range(0,1000) as val CREATE ({v: val})")
+        # Issue a long-running query with a 1-millisecond timeout.
+        try:
+            redis_graph.query("MATCH (a), (b), (c), (d) RETURN *", timeout=1)
+            assert(False)
+        except redis.exceptions.ResponseError as e:
+            assert("Query timed out" in e.args)
+            # Expecting an error.
+            pass
+
+        try:
+            redis_graph.query("RETURN 1", timeout="str")
+            assert(False)
+        except Exception as e:
+            assert("Timeout argument must be a positive integer" in e.args)
+            # Expecting an error.
+            pass
+
 if __name__ == '__main__':
     unittest.main()