Support query timeout (#98)

* Support query timeout

* snapshot repo

Co-authored-by: Guy Korland <[email protected]>

Author: sazzad16

pom.xml

@@ -47,11 +47,18 @@
 		</license>
 	</licenses>
 
+	<!-- temporarily adding jedis snapshot repository -->
+        <repositories>
+            <repository>
+                <id>snapshots-repo</id>
+                <url>https://oss.sonatype.org/content/repositories/snapshots</url>
+            </repository>
+        </repositories>
 	<dependencies>
 		<dependency>
 			<groupId>redis.clients</groupId>
 			<artifactId>jedis</artifactId>
-			<version>3.4.1</version>
+			<version>3.5.0-SNAPSHOT</version>
 		</dependency>
 		<dependency>
 		    <groupId>org.apache.commons</groupId>

src/main/java/com/redislabs/redisgraph/RedisGraph.java

@@ -14,6 +14,14 @@ public interface RedisGraph extends Closeable {
      */
     ResultSet query(String graphId, String query);
 
+    /**
+     * Execute a Cypher query with timeout.
+     * @param graphId a graph to perform the query on
+     * @param query Cypher query
+     * @param timeout
+     * @return a result set
+     */
+    ResultSet query(String graphId, String query, long timeout);
 
     /**
      * Execute a Cypher query with arguments
@@ -35,6 +43,16 @@ public interface RedisGraph extends Closeable {
      */
     ResultSet query(String graphId, String query, Map<String, Object> params);
 
+    /**
+     * Executes a cypher query with parameters and timeout.
+     * @param graphId a graph to perform the query on.
+     * @param query Cypher query.
+     * @param params parameters map.
+     * @param timeout
+     * @return a result set.
+     */
+    ResultSet query(String graphId, String query, Map<String, Object> params, long timeout);
+
     /**
      * Invokes stored procedures without arguments
      * @param graphId a graph to perform the query on

src/main/java/com/redislabs/redisgraph/RedisGraphTransaction.java

@@ -22,6 +22,7 @@ public interface RedisGraphTransaction extends
         MultiKeyCommandsPipeline, ClusterPipeline,
         BinaryScriptingCommandsPipeline, ScriptingCommandsPipeline,
         BasicRedisPipeline, BinaryRedisPipeline, RedisPipeline, Closeable {
+
     /**
      * Execute a Cypher query.
      * @param graphId a graph to perform the query on
@@ -30,6 +31,15 @@ public interface RedisGraphTransaction extends
      */
     Response<ResultSet> query(String graphId, String query);
 
+    /**
+     * Execute a Cypher query with timeout.
+     * @param graphId a graph to perform the query on
+     * @param query Cypher query
+     * @param timeout
+     * @return a response which builds the result set with the query answer.
+     */
+    Response<ResultSet> query(String graphId, String query, long timeout);
+
     /**
      * Execute a Cypher query with arguments
      * @param graphId a graph to perform the query on
@@ -50,6 +60,16 @@ public interface RedisGraphTransaction extends
      */
     Response<ResultSet> query(String graphId, String query, Map<String, Object> params);
 
+    /**
+     * Executes a cypher query with parameters and timeout.
+     * @param graphId a graph to perform the query on.
+     * @param query Cypher query.
+     * @param params parameters map.
+     * @param timeout
+     * @return  a response which builds the result set with the query answer.
+     */
+    Response<ResultSet> query(String graphId, String query, Map<String, Object> params, long timeout);
+
     /**
      * Invokes stored procedures without arguments
      * @param graphId a graph to perform the query on

src/main/java/com/redislabs/redisgraph/impl/Utils.java

@@ -14,6 +14,7 @@ public class Utils {
     public static final List<String> DUMMY_LIST = new ArrayList<>(0);
     public static final Map<String, List<String>> DUMMY_MAP = new HashMap<>(0);
     public static final String COMPACT_STRING = "--COMPACT";
+    public static final String TIMEOUT_STRING = "TIMEOUT";
 
     private static final CharSequenceTranslator ESCAPE_CHYPER;
     static {

src/main/java/com/redislabs/redisgraph/impl/api/AbstractRedisGraph.java

@@ -27,6 +27,15 @@ public abstract class AbstractRedisGraph implements RedisGraph {
      */
     protected abstract ResultSet sendQuery(String graphId, String preparedQuery);
 
+    /**
+     * Sends a query to the redis graph.Implementation and context dependent
+     * @param graphId graph to be queried
+     * @param preparedQuery prepared query
+     * @param timeout
+     * @return Result set
+     */
+    protected abstract ResultSet sendQuery(String graphId, String preparedQuery, long timeout);
+
     /**
      * Execute a Cypher query.
      * @param graphId a graph to perform the query on
@@ -37,6 +46,18 @@ public ResultSet query(String graphId, String query) {
         return sendQuery(graphId, query);
     }
 
+    /**
+     * Execute a Cypher query with timeout.
+     * @param graphId a graph to perform the query on
+     * @param timeout
+     * @param query Cypher query
+     * @return a result set
+     */
+    @Override
+    public ResultSet query(String graphId, String query, long timeout) {
+        return sendQuery(graphId, query, timeout);
+    }
+
     /**
      * Execute a Cypher query with arguments
      * @param graphId a graph to perform the query on
@@ -63,6 +84,19 @@ public ResultSet query(String graphId, String query, Map<String, Object> params)
         return sendQuery(graphId, preparedQuery);
     }
 
+    /**
+     * Executes a cypher query with parameters and timeout.
+     * @param graphId a graph to perform the query on.
+     * @param timeout
+     * @param query Cypher query.
+     * @param params parameters map.
+     * @return a result set.
+     */
+    @Override
+    public ResultSet query(String graphId, String query, Map<String, Object> params, long timeout) {
+        String preparedQuery = Utils.prepareQuery(query, params);
+        return sendQuery(graphId, preparedQuery, timeout);
+    }
 
     public ResultSet callProcedure(String graphId, String procedure){
         return callProcedure(graphId, procedure, Utils.DUMMY_LIST, Utils.DUMMY_MAP);

src/main/java/com/redislabs/redisgraph/impl/api/ContextedRedisGraph.java

@@ -60,6 +60,29 @@ protected ResultSet sendQuery(String graphId, String preparedQuery) {
         }
     }
 
+    /**
+     * Sends the query over the instance only connection
+     * @param graphId graph to be queried
+     * @param timeout
+     * @param preparedQuery prepared query
+     * @return Result set with the query answer
+     */
+    @Override
+    protected ResultSet sendQuery(String graphId, String preparedQuery, long timeout) {
+        Jedis conn = getConnection();
+        try {
+            List<Object> rawResponse = (List<Object>) conn.sendBlockingCommand(RedisGraphCommand.QUERY,
+                graphId, preparedQuery, Utils.COMPACT_STRING, Utils.TIMEOUT_STRING, Long.toString(timeout));
+            return new ResultSetImpl(rawResponse, this, caches.getGraphCache(graphId));
+        }
+        catch (JRedisGraphRunTimeException rt) {
+            throw rt;
+        }
+        catch (JedisDataException j) {
+            throw new JRedisGraphCompileTimeException(j);
+        }
+    }
+
     /**
      * @return Returns the instance Jedis connection.
      */

src/main/java/com/redislabs/redisgraph/impl/api/RedisGraph.java

@@ -69,6 +69,21 @@ protected ResultSet sendQuery(String graphId, String preparedQuery){
         }
     }
 
+    /**
+     * Overrides the abstract function.
+     * Sends the query from any Jedis connection received from the Jedis pool and closes it once done
+     * @param graphId graph to be queried
+     * @param preparedQuery prepared query
+     * @param timeout
+     * @return Result set with the query answer
+     */
+    @Override
+    protected ResultSet sendQuery(String graphId, String preparedQuery, long timeout){
+        try (ContextedRedisGraph contextedRedisGraph = new ContextedRedisGraph(getConnection())) {
+            contextedRedisGraph.setRedisGraphCaches(caches);
+            return contextedRedisGraph.sendQuery(graphId, preparedQuery, timeout);
+        }
+    }
 
     /**
      * Closes the Jedis pool

src/main/java/com/redislabs/redisgraph/impl/api/RedisGraphTransaction.java

@@ -47,6 +47,26 @@ public ResultSet build(Object o) {
         });
     }
 
+    /**
+     * Execute a Cypher query with timeout.
+     *
+     * NOTE: timeout is simply sent to DB. Socket timeout will not be changed.
+     * @param graphId a graph to perform the query on
+     * @param query Cypher query
+     * @param timeout
+     * @return a response which builds the result set with the query answer.
+     */
+    @Override
+    public Response<ResultSet> query(String graphId, String query, long timeout) {
+        client.sendCommand(RedisGraphCommand.QUERY, graphId, query, "--COMPACT", "TIMEOUT", Long.toString(timeout));
+        return getResponse(new Builder<ResultSet>() {
+            @Override
+            public ResultSet build(Object o) {
+                return new ResultSetImpl((List<Object>)o, redisGraph, caches.getGraphCache(graphId));
+            }
+        });
+    }
+
     /**
      * Execute a Cypher query with arguments
      *
@@ -88,6 +108,28 @@ public ResultSet build(Object o) {
         });
     }
 
+    /**
+     * Executes a cypher query with parameters and timeout.
+     *
+     * NOTE: timeout is simply sent to DB. Socket timeout will not be changed.
+     * timeout.
+     * @param graphId a graph to perform the query on.
+     * @param query Cypher query.
+     * @param params parameters map.
+     * @param timeout
+     * @return  a response which builds the result set with the query answer.
+     */
+    @Override
+    public Response<ResultSet> query(String graphId, String query, Map<String, Object> params, long timeout) {
+        String preparedQuery = Utils.prepareQuery(query, params);
+        client.sendCommand(RedisGraphCommand.QUERY, graphId, preparedQuery, "--COMPACT", "TIMEOUT", Long.toString(timeout));
+        return getResponse(new Builder<ResultSet>() {
+            @Override
+            public ResultSet build(Object o) {
+                return new ResultSetImpl((List<Object>)o, redisGraph, caches.getGraphCache(graphId));
+            }
+        });
+    }
 
     /**
      * Invokes stored procedures without arguments, in multi/exec context

src/test/java/com/redislabs/redisgraph/RedisGraphAPITest.java

@@ -1,22 +1,20 @@
 package com.redislabs.redisgraph;
 
-
-import java.util.*;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
-
+import com.redislabs.redisgraph.Statistics.Label;
 import com.redislabs.redisgraph.graph_entities.Edge;
 import com.redislabs.redisgraph.graph_entities.Node;
 import com.redislabs.redisgraph.graph_entities.Path;
 import com.redislabs.redisgraph.graph_entities.Property;
 import com.redislabs.redisgraph.impl.api.RedisGraph;
+import com.redislabs.redisgraph.impl.resultset.RecordImpl;
 import com.redislabs.redisgraph.impl.resultset.ResultSetImpl;
 import com.redislabs.redisgraph.test.utils.PathBuilder;
-import org.junit.*;
+import redis.clients.jedis.exceptions.JedisDataException;
 
-import com.redislabs.redisgraph.Statistics.Label;
-
-import static com.redislabs.redisgraph.Header.ResultSetColumnTypes.*;
+import java.util.*;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.junit.*;
 
 public class RedisGraphAPITest {
 
@@ -29,20 +27,18 @@ public RedisGraphAPITest() {
     public void createApi(){
         api = new RedisGraph();
     }
+
     @After
     public void deleteGraph() {
-
         api.deleteGraph("social");
         api.close();
     }
 
-
     @Test
     public void testCreateNode() {
         // Create a node    	
         ResultSet resultSet = api.query("social", "CREATE ({name:'roi',age:32})");
 
-
         Assert.assertEquals(1, resultSet.getStatistics().nodesCreated());
         Assert.assertNull(resultSet.getStatistics().getStringValue(Label.NODES_DELETED));
         Assert.assertNull(resultSet.getStatistics().getStringValue(Label.RELATIONSHIPS_CREATED));
@@ -116,8 +112,6 @@ public void testDeleteNodes(){
 
         Assert.assertNotNull(deleteResult.getStatistics().getStringValue(Label.QUERY_INTERNAL_EXECUTION_TIME));
 
-
-
     }
 
     @Test
@@ -192,7 +186,6 @@ public void testHeader(){
         Assert.assertEquals("r", schemaNames.get(1));
         Assert.assertEquals("a.age", schemaNames.get(2));
 
-
     }
 
     @Test
@@ -311,7 +304,6 @@ public void testRecord(){
 
     }
 
-
     @Ignore
     @Test
     public void tinyTestMultiThread(){
@@ -322,9 +314,7 @@ public void tinyTestMultiThread(){
                     mapToObj(
                             j-> api.query("social", "MATCH (a:person) RETURN a")).
                     collect(Collectors.toList());
-
         }
-
     }
 
     @Test
@@ -484,15 +474,13 @@ record = resultSet.next();
 
     }
 
-
     @Test
     public void testEscapedQuery() {
         Assert.assertNotNull(api.query("social", "CREATE (:escaped{s1:%s,s2:%s})", "S\"'", "S'\""));
         Assert.assertNotNull(api.query("social", "MATCH (n) where n.s1=%s and n.s2=%s RETURN n", "S\"'", "S'\""));
         Assert.assertNotNull(api.query("social", "MATCH (n) where n.s1='S\"' RETURN n"));
     }
 
-
     @Test
     public void testMultiExec(){
         try (RedisGraphContext c = api.getContext()) {
@@ -965,4 +953,12 @@ public void testCachedExecution() {
         Assert.assertEquals(params.get("val"), r.getValue(0));
         Assert.assertTrue(resultSet.getStatistics().cachedExecution());
     }
+
+    @Test
+    public void timeoutArgument() {
+        ResultSet rs = api.query("social", "UNWIND range(0,100) AS x WITH x AS x WHERE x = 100 RETURN x", 1L);
+        Assert.assertEquals(1, rs.size());
+        Record r = rs.next();
+        Assert.assertEquals(Long.valueOf(100), r.getValue(0));
+    }
 }

src/test/java/com/redislabs/redisgraph/exceptions/JRedisGraphErrorTest.java

@@ -138,4 +138,12 @@ public void testContextExceptionFlow() {
         }
 
     }
+
+    @Test
+    public void timeoutExcpetion() {
+        exceptionRule.expect(JRedisGraphRunTimeException.class);
+        exceptionRule.expectMessage("Query timed out");
+
+        api.query("social", "UNWIND range(0,100000) AS x WITH x AS x WHERE x = 10000 RETURN x", 1L);
+    }
 }