Remove duplicate code for query processing

Author: mangalaman93

dgraph/cmd/alpha/http.go

@@ -21,7 +21,6 @@ import (
 	"compress/gzip"
 	"context"
 	"encoding/json"
-	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -158,6 +157,26 @@ func writeResponse(w http.ResponseWriter, r *http.Request, b []byte) (int, error
 	return out.Write(b)
 }
 
+func writeEntry(out *bytes.Buffer, key string, js []byte) error {
+	if _, err := out.WriteRune('"'); err != nil {
+		return err
+	}
+	if _, err := out.WriteString(key); err != nil {
+		return err
+	}
+	if _, err := out.WriteRune('"'); err != nil {
+		return err
+	}
+	if _, err := out.WriteRune(':'); err != nil {
+		return err
+	}
+	if _, err := out.Write(js); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // This method should just build the request and proxy it to the Query method of dgraph.Server.
 // It can then encode the response as appropriate before sending it back to the user.
 func queryHandler(w http.ResponseWriter, r *http.Request) {
@@ -254,25 +273,6 @@ func queryHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	var out bytes.Buffer
-	writeEntry := func(key string, js []byte) error {
-		if _, err := out.WriteRune('"'); err != nil {
-			return err
-		}
-		if _, err := out.WriteString(key); err != nil {
-			return err
-		}
-		if _, err := out.WriteRune('"'); err != nil {
-			return err
-		}
-		if _, err := out.WriteRune(':'); err != nil {
-			return err
-		}
-		if _, err := out.Write(js); err != nil {
-			return err
-		}
-		return nil
-	}
-
 	e := query.Extensions{
 		Txn:     resp.Txn,
 		Latency: resp.Latency,
@@ -288,15 +288,15 @@ func queryHandler(w http.ResponseWriter, r *http.Request) {
 		x.SetStatusWithData(w, x.Error, err.Error())
 		return
 	}
-	if err := writeEntry("data", resp.Json); err != nil {
+	if err := writeEntry(&out, "data", resp.Json); err != nil {
 		x.SetStatusWithData(w, x.Error, err.Error())
 		return
 	}
 	if _, err := out.WriteRune(','); err != nil {
 		x.SetStatusWithData(w, x.Error, err.Error())
 		return
 	}
-	if err := writeEntry("extensions", js); err != nil {
+	if err := writeEntry(&out, "extensions", js); err != nil {
 		x.SetStatusWithData(w, x.Error, err.Error())
 		return
 	}
@@ -409,22 +409,61 @@ func mutationHandler(w http.ResponseWriter, r *http.Request) {
 		e.Txn.Keys = e.Txn.Keys[:0]
 	}
 
-	response := map[string]interface{}{}
-	response["extensions"] = e
-	mp := map[string]interface{}{}
-	mp["code"] = x.Success
-	mp["message"] = "Done"
-	mp["uids"] = resp.Uids
-	if len(resp.Vars) > 0 {
-		vars := make(map[string][]string)
-		// Flatten the mutated map so that it is easier to parse for the client.
-		for v, uids := range resp.Vars {
-			vars[fmt.Sprintf("uid(%s)", v)] = uids.GetUids()
+	var out bytes.Buffer
+	// If response is either empty or only contains empty braces {}, it's a no-op
+	if len(resp.Json) > 2 {
+		var i int
+		for i = len(resp.Json) - 1; i >= 0; i-- {
+			if resp.Json[i] == '}' {
+				break
+			}
+		}
+
+		out = *bytes.NewBuffer(resp.Json[:i])
+		if _, err := out.WriteRune(','); err != nil {
+			x.SetStatusWithData(w, x.Error, err.Error())
+			return
+		}
+	} else {
+		if _, err := out.WriteRune('{'); err != nil {
+			x.SetStatusWithData(w, x.Error, err.Error())
+			return
 		}
-		mp["vars"] = vars
 	}
-	response["data"] = mp
 
+	if err := writeEntry(&out, "code", []byte("\""+x.Success+"\"")); err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+	if _, err := out.WriteRune(','); err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+	if err := writeEntry(&out, "message", []byte("\"Done\"")); err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+	if _, err := out.WriteRune(','); err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+	uids, err := json.Marshal(resp.Uids)
+	if err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+	if err := writeEntry(&out, "uids", uids); err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+	if _, err := out.WriteRune('}'); err != nil {
+		x.SetStatusWithData(w, x.Error, err.Error())
+		return
+	}
+
+	response := map[string]interface{}{}
+	response["extensions"] = e
+	response["data"] = json.RawMessage(out.Bytes())
 	js, err := json.Marshal(response)
 	if err != nil {
 		x.SetStatusWithData(w, x.Error, err.Error())

dgraph/cmd/alpha/http_test.go

@@ -166,7 +166,7 @@ type mutationResponse struct {
 	keys    []string
 	preds   []string
 	startTs uint64
-	vars    map[string][]string
+	data    json.RawMessage
 }
 
 func mutationWithTs(m, t string, isJson bool, commitNow bool, ts uint64) (
@@ -193,18 +193,22 @@ func mutationWithTs(m, t string, isJson bool, commitNow bool, ts uint64) (
 		return mr, err
 	}
 
-	type resData struct {
-		MutationVars map[string][]string `json:"vars"`
-	}
-	var rd resData
-	if err := json.Unmarshal(r.Data, &rd); err != nil {
-		return mr, err
-	}
-
-	mr.vars = rd.MutationVars
 	mr.keys = r.Extensions.Txn.Keys
 	mr.preds = r.Extensions.Txn.Preds
 	mr.startTs = r.Extensions.Txn.StartTs
+	sort.Strings(mr.preds)
+
+	var d map[string]interface{}
+	if err := json.Unmarshal(r.Data, &d); err != nil {
+		return mr, err
+	}
+	delete(d, "code")
+	delete(d, "message")
+	delete(d, "uids")
+	mr.data, err = json.Marshal(d)
+	if err != nil {
+		return mr, err
+	}
 	return mr, nil
 }