Add optional timeout argument to graph queries

README.md

@@ -123,6 +123,17 @@ Name: John Doe
 Age: 33
 ```
 
+## Running queries with timeouts
+
+Queries can be run with a millisecond-level timeout as described in [the module documentation](https://oss.redis.com/redisgraph/configuration/#timeout). To take advantage of this feature, the `QueryOptions` struct should be used:
+
+```go
+options := NewQueryOptions().SetTimeout(10) // 10-millisecond timeout
+res, err := graph.QueryWithOptions("MATCH (src {name: 'John Doe'})-[*]->(dest) RETURN dest", options)
+```
+
+`ParameterizedQueryWithOptions` and `ROQueryWithOptions` endpoints are also exposed by the client.
+
 ## Running tests
 
 A simple test suite is provided, and can be run with:

client_test.go

@@ -473,3 +473,22 @@ func TestNodeMapDatatype(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, 1, res.RelationshipsDeleted(), "Expecting 1 relationships deleted")
 }
+
+func TestTimeout(t *testing.T) {
+	// Instantiate a new QueryOptions struct with a 1-second timeout
+	options := NewQueryOptions().SetTimeout(1)
+
+	// Verify that the timeout was set properly
+	assert.Equal(t, 1, options.GetTimeout())
+
+	// Issue a long-running query with a 1-millisecond timeout.
+	res, err := graph.QueryWithOptions("UNWIND range(0, 1000000) AS v RETURN v", options)
+	assert.Nil(t, res)
+	assert.NotNil(t, err)
+
+	params := make(map[string]interface{})
+	params["ub"] = 1000000
+	res, err = graph.ParameterizedQueryWithOptions("UNWIND range(0, $ub) AS v RETURN v", params, options);
+	assert.Nil(t, res)
+	assert.NotNil(t, err)
+}

graph.go

@@ -8,6 +8,11 @@ import (
 	"github.com/gomodule/redigo/redis"
 )
 
+// QueryOptions are a set of additional arguments to be emitted with a query.
+type QueryOptions struct {
+	timeout           int
+}
+
 // Graph represents a graph, which is a collection of nodes and edges.
 type Graph struct {
 	Id                string
@@ -97,9 +102,26 @@ func (g *Graph) Commit() (*QueryResult, error) {
 	return g.Query(q)
 }
 
+// NewQueryOptions instantiates a new QueryOptions struct.
+func NewQueryOptions() *QueryOptions {
+	return &QueryOptions{
+		timeout:               -1,
+	}
+}
+
+// SetTimeout sets the timeout member of the QueryOptions struct
+func (options *QueryOptions) SetTimeout(timeout int) *QueryOptions {
+	options.timeout = timeout
+	return options
+}
+
+// GetTimeout retrieves the timeout of the QueryOptions struct
+func (options *QueryOptions) GetTimeout() int {
+	return options.timeout
+}
+
 // Query executes a query against the graph.
 func (g *Graph) Query(q string) (*QueryResult, error) {
-
 	r, err := g.Conn.Do("GRAPH.QUERY", g.Id, q, "--compact")
 	if err != nil {
 		return nil, err
@@ -126,6 +148,46 @@ func (g *Graph) ParameterizedQuery(q string, params map[string]interface{}) (*Qu
 	return g.Query(q);
 }
 
+// QueryWithOptions issues a query with the given timeout
+func (g *Graph) QueryWithOptions(q string, options *QueryOptions) (*QueryResult, error) {
+	var r interface{}
+	var err error
+	if(options.timeout >= 0) {
+		r, err = g.Conn.Do("GRAPH.QUERY", g.Id, q, "--compact", "timeout", options.timeout)
+	} else {
+		r, err = g.Conn.Do("GRAPH.QUERY", g.Id, q, "--compact")
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	return QueryResultNew(g, r)
+}
+
+// ParameterizedQueryWithOptions issues a parameterized query with the given timeout
+func (g *Graph) ParameterizedQueryWithOptions(q string, params map[string]interface{}, options *QueryOptions) (*QueryResult, error) {
+	if(params != nil){
+		q = BuildParamsHeader(params) + q
+	}
+	return g.QueryWithOptions(q, options);
+}
+
+// ROQueryWithOptions issues a read-only query with the given timeout
+func (g *Graph) ROQueryWithOptions(q string, options *QueryOptions) (*QueryResult, error) {
+	var r interface{}
+	var err error
+	if(options.timeout >= 0) {
+		r, err = g.Conn.Do("GRAPH.RO_QUERY", g.Id, q, "--compact", "timeout", options.timeout)
+	} else {
+		r, err = g.Conn.Do("GRAPH.RO_QUERY", g.Id, q, "--compact")
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	return QueryResultNew(g, r)
+}
+
 // Merge pattern
 func (g *Graph) Merge(p string) (*QueryResult, error) {
 	q := fmt.Sprintf("MERGE %s", p)