Merge pull request #23 from boatilus/support-count

Add support for returning row count

.gitignore

@@ -18,4 +18,9 @@
 .idea
 
 #VSCode Folder Ignore
-.vscode
+.vscode
+
+# Supabase
+supabase
+
+.env

execute.go

@@ -7,8 +7,12 @@ import (
 	"io"
 	"net/http"
 	"path"
+	"strconv"
+	"strings"
 )
 
+type countType = int64
+
 // ExecuteError is the error response format from postgrest. We really
 // only use Code and Message, but we'll keep it as a struct for now.
 
@@ -19,16 +23,16 @@ 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, error) {
+func executeHelper(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, client.ClientError
+		return nil, 0, client.ClientError
 	}
 
 	readerBody := bytes.NewBuffer(body)
 	baseUrl := path.Join(append([]string{client.clientTransport.baseURL.Path}, urlFragments...)...)
 	req, err := http.NewRequest(method, baseUrl, readerBody)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
 	for key, val := range headers {
@@ -41,42 +45,56 @@ func executeHelper(client *Client, method string, body []byte, urlFragments []st
 	req.URL.RawQuery = q.Encode()
 	resp, err := client.session.Do(req)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
 	respbody, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
 	// https://postgrest.org/en/stable/api.html#errors-and-http-status-codes
 	if resp.StatusCode >= 400 {
 		var errmsg *ExecuteError
 		err := json.Unmarshal(respbody, &errmsg)
 		if err != nil {
-			return nil, err
+			return nil, 0, err
+		}
+		return nil, 0, fmt.Errorf("(%s) %s", errmsg.Code, errmsg.Message)
+	}
+
+	var count countType
+
+	contentRange := resp.Header.Get("Content-Range")
+	if contentRange != "" {
+		split := strings.Split(contentRange, "/")
+		if len(split) > 1 && split[1] != "*" {
+			count, err = strconv.ParseInt(split[1], 0, 64)
+			if err != nil {
+				return nil, 0, err
+			}
 		}
-		return nil, fmt.Errorf("(%s) %s", errmsg.Code, errmsg.Message)
 	}
 
 	err = resp.Body.Close()
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
-	return respbody, nil
+
+	return respbody, count, nil
 }
 
-func executeString(client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) (string, error) {
-	resp, err := executeHelper(client, method, body, urlFragments, headers, params)
-	return string(resp), err
+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)
+	return string(resp), count, err
 }
 
-func execute(client *Client, method string, body []byte, urlFragments []string, headers map[string]string, params map[string]string) ([]byte, error) {
+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 executeTo(client *Client, method string, body []byte, to interface{}, urlFragments []string, headers map[string]string, params map[string]string) error {
-	resp, err := executeHelper(client, method, body, urlFragments, headers, params)
+	resp, _, err := executeHelper(client, method, body, urlFragments, headers, params)
 
 	if err != nil {
 		return err

export_test.go

@@ -0,0 +1,55 @@
+package postgrest
+
+import (
+	"os"
+	"regexp"
+	"testing"
+)
+
+const urlEnv = "POSTGREST_URL"
+const apiKeyEnv = "API_KEY"
+
+// If false, mock responses with httpmock. If true, use POSTGREST_URL (and
+// optionally, API_KEY for Supabase), to run tests against an actual Postgres
+// instance.
+var mockResponses bool = false
+
+var mockPath *regexp.Regexp
+
+// A mock table/result set.
+var users = []map[string]interface{}{
+	{
+		"id":    float64(1), // numeric types are returned as float64s
+		"name":  "sean",
+		"email": "sean@test.com",
+	},
+	{
+		"id":    float64(2),
+		"name":  "patti",
+		"email": "patti@test.com",
+	},
+}
+
+func createClient(t *testing.T) *Client {
+	// If a POSTGREST_URL environment variable is specified, we'll use that
+	// to test against real endpoints.
+	url := os.Getenv(urlEnv)
+	if url == "" {
+		url = "http://mock.xyz"
+		mockResponses = true
+
+		var err error
+		mockPath, err = regexp.Compile(regexp.QuoteMeta(url) + "?.*")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	headers := make(map[string]string)
+	if apiKeyEnv != "" {
+		// If the API_KEY env is specified, we'll use it to auth with Supabase.
+		headers["apikey"] = os.Getenv(apiKeyEnv)
+	}
+
+	return NewClient(url, "", headers)
+}

filterbuilder.go

@@ -16,11 +16,11 @@ type FilterBuilder struct {
 	params    map[string]string
 }
 
-func (f *FilterBuilder) ExecuteString() (string, error) {
+func (f *FilterBuilder) ExecuteString() (string, countType, error) {
 	return executeString(f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
 }
 
-func (f *FilterBuilder) Execute() ([]byte, error) {
+func (f *FilterBuilder) Execute() ([]byte, countType, error) {
 	return execute(f.client, f.method, f.body, []string{f.tableName}, f.headers, f.params)
 }
 

go.mod

@@ -1,3 +1,8 @@
 module github.com/supabase/postgrest-go
 
 go 1.16
+
+require (
+	github.com/jarcoal/httpmock v1.1.0
+	github.com/stretchr/testify v1.7.0
+)

go.sum

@@ -0,0 +1,13 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/jarcoal/httpmock v1.1.0 h1:F47ChZj1Y2zFsCXxNkBPwNNKnAyOATcdQibk0qEdVCE=
+github.com/jarcoal/httpmock v1.1.0/go.mod h1:ATjnClrvW/3tijVmpL/va5Z3aAyGvqU3gCT8nX0Txik=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

querybuilder.go

@@ -15,11 +15,11 @@ type QueryBuilder struct {
 	params    map[string]string
 }
 
-func (q *QueryBuilder) ExecuteString() (string, error) {
+func (q *QueryBuilder) ExecuteString() (string, countType, error) {
 	return executeString(q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
 }
 
-func (q *QueryBuilder) Execute() ([]byte, error) {
+func (q *QueryBuilder) Execute() ([]byte, countType, error) {
 	return execute(q.client, q.method, q.body, []string{q.tableName}, q.headers, q.params)
 }
 
@@ -53,11 +53,12 @@ func (q *QueryBuilder) Select(columns, count string, head bool) *FilterBuilder {
 	}
 
 	if count != "" && (count == `exact` || count == `planned` || count == `estimated`) {
-		currentValue, ok := q.params["Prefer"]
+		currentValue, ok := q.headers["Prefer"]
 		if ok && currentValue != "" {
-			count = fmt.Sprintf("%s,count=%s", currentValue, count)
+			q.headers["Prefer"] = fmt.Sprintf("%s,count=%s", currentValue, count)
+		} else {
+			q.headers["Prefer"] = fmt.Sprintf("count=%s", count)
 		}
-		q.params["Prefer"] = count
 	}
 	return &FilterBuilder{client: q.client, method: q.method, body: q.body, tableName: q.tableName, headers: q.headers, params: q.params}
 }

querybuilder_test.go

@@ -0,0 +1,83 @@
+package postgrest
+
+import (
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/jarcoal/httpmock"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewClient(t *testing.T) {
+	assert.NotNil(t, NewClient("", "", nil))
+}
+
+func TestSelect(t *testing.T) {
+	assert := assert.New(t)
+	c := createClient(t)
+
+	t.Run("ValidResult", func(t *testing.T) {
+		got := []map[string]interface{}{}
+
+		if mockResponses {
+			httpmock.Activate()
+			defer httpmock.DeactivateAndReset()
+
+			responder, _ := httpmock.NewJsonResponder(200, users)
+			httpmock.RegisterRegexpResponder("GET", mockPath, responder)
+		}
+
+		bs, count, err := c.From("users").Select("id, name, email", "", false).Execute()
+		assert.NoError(err)
+
+		err = json.Unmarshal(bs, &got)
+		assert.NoError(err)
+		assert.EqualValues(users, got)
+		assert.Equal(countType(0), count)
+	})
+
+	t.Run("WithCount", func(t *testing.T) {
+		got := []map[string]interface{}{}
+
+		if mockResponses {
+			httpmock.Activate()
+			defer httpmock.DeactivateAndReset()
+
+			httpmock.RegisterRegexpResponder("GET", mockPath, func(req *http.Request) (*http.Response, error) {
+				resp, _ := httpmock.NewJsonResponse(200, users)
+
+				resp.Header.Add("Content-Range", "0-1/2")
+				return resp, nil
+			})
+		}
+
+		bs, count, err := c.From("users").Select("id, name, email", "exact", false).Execute()
+		assert.NoError(err)
+
+		err = json.Unmarshal(bs, &got)
+		assert.NoError(err)
+		assert.EqualValues(users, got)
+		assert.Equal(countType(2), count)
+	})
+}
+
+func TestFilter(t *testing.T) {
+	assert := assert.New(t)
+	c := createClient(t)
+
+	t.Run("Eq", func(t *testing.T) {
+		want := "[{\"email\":\"patti@test.com\"}]"
+
+		if mockResponses {
+			httpmock.Activate()
+			defer httpmock.DeactivateAndReset()
+
+			httpmock.RegisterRegexpResponder("GET", mockPath, httpmock.NewStringResponder(200, want))
+		}
+
+		got, _, err := c.From("users").Select("email", "", false).Eq("email", "patti@test.com").ExecuteString()
+		assert.NoError(err)
+		assert.Equal(want, got)
+	})
+}

test/basic/main.go

@@ -19,10 +19,11 @@ var (
 func main() {
 	client := postgrest.NewClient(REST_URL, schema, headers)
 
-	res, err := client.From("todos").Select("id,task,done", "", false).Eq("task", "that created from postgrest-go").Execute()
+	res, count, err := client.From("todos").Select("id,task,done", "", false).Eq("task", "that created from postgrest-go").Execute()
 	if err != nil {
 		panic(err)
 	}
 
 	fmt.Println(res)
+	fmt.Printf("count: %v", count)
 }

transformbuilder.go

@@ -13,11 +13,11 @@ type TransformBuilder struct {
 	params  map[string]string
 }
 
-func (t *TransformBuilder) ExecuteString() (string, error) {
+func (t *TransformBuilder) ExecuteString() (string, countType, error) {
 	return executeString(t.client, t.method, t.body, []string{}, t.headers, t.params)
 }
 
-func (t *TransformBuilder) Execute() ([]byte, error) {
+func (t *TransformBuilder) Execute() ([]byte, countType, error) {
 	return execute(t.client, t.method, t.body, []string{}, t.headers, t.params)
 }