Merge pull request #49 from assembly-winston/addContexts

Added ctx arguments to execute

execute.go

@@ -2,6 +2,7 @@ package postgrest
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -25,14 +26,14 @@ type ExecuteError struct {
 	Message string `json:"message"`
 }
 
-func executeHelper(client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) ([]byte, countType, error) {
+func executeHelper(ctx context.Context, client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) ([]byte, countType, error) {
 	if client.ClientError != nil {
 		return nil, 0, client.ClientError
 	}
 
 	readerBody := bytes.NewBuffer(body)
 	baseUrl := path.Join(append([]string{client.Transport.baseURL.Path}, urlFragments...)...)
-	req, err := http.NewRequest(method, baseUrl, readerBody)
+	req, err := http.NewRequestWithContext(ctx, method, baseUrl, readerBody)
 	if err != nil {
 		return nil, 0, fmt.Errorf("error creating request: %s", err.Error())
 	}
@@ -86,17 +87,17 @@ func executeHelper(client *Client, method string, body []byte, urlFragments []st
 	return respBody, count, nil
 }
 
-func executeString(client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) (string, countType, error) {
-	resp, count, err := executeHelper(client, method, body, urlFragments, headers, params)
+func executeString(ctx context.Context, client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) (string, countType, error) {
+	resp, count, err := executeHelper(ctx, client, method, body, urlFragments, headers, params)
 	return string(resp), count, err
 }
 
-func execute(client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) ([]byte, countType, error) {
-	return executeHelper(client, method, body, urlFragments, headers, params)
+func execute(ctx context.Context, client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) ([]byte, countType, error) {
+	return executeHelper(ctx, client, method, body, urlFragments, headers, params)
 }
 
-func executeTo(client *Client, method string, body []byte, to interface{}, urlFragments []string, headers map[string]string, params map[string]string) (countType, error) {
-	resp, count, err := executeHelper(client, method, body, urlFragments, headers, params)
+func executeTo(ctx context.Context, client *Client, method string, body []byte, to interface{}, urlFragments []string, headers map[string]string, params map[string]string) (countType, error) {
+	resp, count, err := executeHelper(ctx, client, method, body, urlFragments, headers, params)
 
 	if err != nil {
 		return count, err

filterbuilder.go

@@ -1,6 +1,7 @@
 package postgrest
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"regexp"
@@ -21,19 +22,38 @@ type FilterBuilder struct {
 // ExecuteString runs the PostgREST query, returning the result as a JSON
 // string.
 func (f *FilterBuilder) ExecuteString() (string, int64, error) {
-	return executeString(f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
+	return executeString(context.Background(), f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
+}
+
+// ExecuteStringWithContext runs the PostgREST query, returning the result as
+// a JSON string.
+func (f *FilterBuilder) ExecuteStringWithContext(ctx context.Context) (string, int64, error) {
+	return executeString(ctx, f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
 }
 
 // Execute runs the PostgREST query, returning the result as a byte slice.
 func (f *FilterBuilder) Execute() ([]byte, int64, error) {
-	return execute(f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
+	return execute(context.Background(), f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
+}
+
+// ExecuteWithContext runs the PostgREST query with the given context,
+// returning the result as a byte slice.
+func (f *FilterBuilder) ExecuteWithContext(ctx context.Context) ([]byte, int64, error) {
+	return execute(ctx, f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
 }
 
 // ExecuteTo runs the PostgREST query, encoding the result to the supplied
 // interface. Note that the argument for the to parameter should always be a
 // reference to a slice.
 func (f *FilterBuilder) ExecuteTo(to interface{}) (countType, error) {
-	return executeTo(f.client, f.method, f.body, to, []string{f.tableName}, f.headers, f.params)
+	return executeTo(context.Background(), f.client, f.method, f.body, to, []string{f.tableName}, f.headers, f.params)
+}
+
+// ExecuteToWithContext runs the PostgREST query with the given context,
+// encoding the result to the supplied interface. Note that the argument for
+// the to parameter should always be a reference to a slice.
+func (f *FilterBuilder) ExecuteToWithContext(ctx context.Context, to interface{}) (countType, error) {
+	return executeTo(ctx, f.client, f.method, f.body, to, []string{f.tableName}, f.headers, f.params)
 }
 
 var filterOperators = []string{"eq", "neq", "gt", "gte", "lt", "lte", "like", "ilike", "is", "in", "cs", "cd", "sl", "sr", "nxl", "nxr", "adj", "ov", "fts", "plfts", "phfts", "wfts"}
@@ -158,7 +178,7 @@ func (f *FilterBuilder) Contains(column string, value []string) *FilterBuilder {
 	}
 
 	valueString := fmt.Sprintf("{%s}", strings.Join(newValue, ","))
-	
+
 	f.params[column] = "cs." + valueString
 	return f
 }
@@ -170,7 +190,7 @@ func (f *FilterBuilder) ContainedBy(column string, value []string) *FilterBuilde
 	}
 
 	valueString := fmt.Sprintf("{%s}", strings.Join(newValue, ","))
-	
+
 	f.params[column] = "cd." + valueString
 	return f
 }

filterbuilder_test.go

@@ -1,10 +1,12 @@
 package postgrest
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"sort"
 	"testing"
+	"time"
 
 	"github.com/jarcoal/httpmock"
 	"github.com/stretchr/testify/assert"
@@ -120,6 +122,25 @@ func TestFilterBuilder_Limit(t *testing.T) {
 	assert.Equal(countType(len(users)), count, "expected count to be %v", len(users))
 }
 
+func TestFilterBuilder_ContextCanceled(t *testing.T) {
+	c := createClient(t)
+	assert := assert.New(t)
+
+	if mockResponses {
+		httpmock.Activate()
+		defer httpmock.DeactivateAndReset()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
+	defer cancel()
+
+	time.Sleep(1 * time.Nanosecond)
+
+	_, _, err := c.From("users").Select("id, name, email", "exact", false).Limit(1, "").ExecuteWithContext(ctx)
+	// This test should immediately fail on a canceled context.
+	assert.Error(err)
+}
+
 func TestFilterBuilder_Order(t *testing.T) {
 	c := createClient(t)
 	assert := assert.New(t)

querybuilder.go

@@ -1,6 +1,7 @@
 package postgrest
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"strings"
@@ -16,22 +17,41 @@ type QueryBuilder struct {
 	params    map[string]string
 }
 
-// ExecuteString runs the Postgrest query, returning the result as a JSON
+// ExecuteString runs the PostgREST query, returning the result as a JSON
 // string.
 func (q *QueryBuilder) ExecuteString() (string, int64, error) {
-	return executeString(q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
+	return executeString(context.Background(), q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
+}
+
+// ExecuteStringWithContext runs the PostgREST query, returning the result as
+// a JSON string.
+func (q *QueryBuilder) ExecuteStringWithContext(ctx context.Context) (string, int64, error) {
+	return executeString(ctx, q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
 }
 
 // Execute runs the Postgrest query, returning the result as a byte slice.
 func (q *QueryBuilder) Execute() ([]byte, int64, error) {
-	return execute(q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
+	return execute(context.Background(), q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
+}
+
+// ExecuteWithContext runs the PostgREST query with the given context,
+// returning the result as a byte slice.
+func (q *QueryBuilder) ExecuteWithContext(ctx context.Context) ([]byte, int64, error) {
+	return execute(ctx, q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
 }
 
-// ExecuteTo runs the Postgrest query, encoding the result to the supplied
+// ExecuteTo runs the PostgREST query, encoding the result to the supplied
 // interface. Note that the argument for the to parameter should always be a
 // reference to a slice.
 func (q *QueryBuilder) ExecuteTo(to interface{}) (int64, error) {
-	return executeTo(q.client, q.method, q.body, to, []string{q.tableName}, q.headers, q.params)
+	return executeTo(context.Background(), q.client, q.method, q.body, to, []string{q.tableName}, q.headers, q.params)
+}
+
+// ExecuteToWithContext runs the PostgREST query with the given context,
+// encoding the result to the supplied interface. Note that the argument for
+// the to parameter should always be a reference to a slice.
+func (q *QueryBuilder) ExecuteToWithContext(ctx context.Context, to interface{}) (int64, error) {
+	return executeTo(ctx, q.client, q.method, q.body, to, []string{q.tableName}, q.headers, q.params)
 }
 
 // Select performs vertical filtering.