Squash mrjn/client branch

Author: Janardhan Reddy

client/client.go

@@ -17,20 +17,14 @@
 package client
 
 import (
-	"context"
 	"fmt"
-	"log"
-	"strings"
-	"sync"
-	"sync/atomic"
 	"time"
 
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-
 	"github.com/dgraph-io/dgraph/protos"
 	"github.com/dgraph-io/dgraph/types"
 	"github.com/dgraph-io/dgraph/x"
+	geom "github.com/twpayne/go-geom"
+	"github.com/twpayne/go-geom/encoding/geojson"
 )
 
 type Op int
@@ -71,6 +65,19 @@ func checkNQuad(nq protos.NQuad) error {
 	return nil
 }
 
+func checkSchema(schema protos.SchemaUpdate) error {
+	typ := types.TypeID(schema.ValueType)
+	if typ == types.UidID && schema.Directive == protos.SchemaUpdate_INDEX {
+		// index on uid type
+		return x.Errorf("Index not allowed on predicate of type uid on predicate %s",
+			schema.Predicate)
+	} else if typ != types.UidID && schema.Directive == protos.SchemaUpdate_REVERSE {
+		// reverse on non-uid type
+		return x.Errorf("Cannot reverse for non-uid type on predicate %s", schema.Predicate)
+	}
+	return nil
+}
+
 // SetQuery sets a query with graphQL variables as part of the request.
 func (req *Req) SetQuery(q string) {
 	req.gr.Query = q
@@ -82,46 +89,31 @@ func (req *Req) SetQueryWithVariables(q string, vars map[string]string) {
 	req.gr.Vars = vars
 }
 
-func (req *Req) addMutation(nq protos.NQuad, op Op) {
+func (req *Req) addMutation(e Edge, op Op) {
 	if req.gr.Mutation == nil {
 		req.gr.Mutation = new(protos.Mutation)
 	}
 
 	if op == SET {
-		req.gr.Mutation.Set = append(req.gr.Mutation.Set, &nq)
+		req.gr.Mutation.Set = append(req.gr.Mutation.Set, &e.nq)
 	} else if op == DEL {
-		req.gr.Mutation.Del = append(req.gr.Mutation.Del, &nq)
+		req.gr.Mutation.Del = append(req.gr.Mutation.Del, &e.nq)
 	}
 }
 
-// AddMutation adds (but does not send) a mutation to the Req object. Mutations
-// are sent when client.Run() is called.
-func (req *Req) AddMutation(nq protos.NQuad, op Op) error {
-	if err := checkNQuad(nq); err != nil {
+func (req *Req) Set(e Edge) error {
+	if err := checkNQuad(e.nq); err != nil {
 		return err
 	}
-	req.addMutation(nq, op)
-	return nil
-}
-
-func AddFacet(key string, val string, nq *protos.NQuad) error {
-	nq.Facets = append(nq.Facets, &protos.Facet{
-		Key: key,
-		Val: val,
-	})
+	req.addMutation(e, SET)
 	return nil
 }
 
-func checkSchema(schema protos.SchemaUpdate) error {
-	typ := types.TypeID(schema.ValueType)
-	if typ == types.UidID && schema.Directive == protos.SchemaUpdate_INDEX {
-		// index on uid type
-		return x.Errorf("Index not allowed on predicate of type uid on predicate %s",
-			schema.Predicate)
-	} else if typ != types.UidID && schema.Directive == protos.SchemaUpdate_REVERSE {
-		// reverse on non-uid type
-		return x.Errorf("Cannot reverse for non-uid type on predicate %s", schema.Predicate)
+func (req *Req) Delete(e Edge) error {
+	if err := checkNQuad(e.nq); err != nil {
+		return err
 	}
+	req.addMutation(e, DEL)
 	return nil
 }
 
@@ -149,159 +141,179 @@ func (req *Req) reset() {
 }
 
 type nquadOp struct {
-	nq protos.NQuad
+	e  Edge
 	op Op
 }
 
-// BatchMutation is used to batch mutations and send them off to the server
-// concurrently. It is useful while doing migrations and bulk data loading.
-// It is possible to control the batch size and the number of concurrent requests
-// to make.
-type BatchMutation struct {
-	size    int
-	pending int
-
-	nquads chan nquadOp
-	schema chan protos.SchemaUpdate
-	dc     protos.DgraphClient
-	wg     sync.WaitGroup
-
-	// Miscellaneous information to print counters.
-	// Num of RDF's sent
-	rdfs uint64
-	// Num of mutations sent
-	mutations uint64
-	// To get time elapsed.
-	start time.Time
+type Node uint64
+
+func (n Node) String() string {
+	return fmt.Sprintf("%#x", uint64(n))
 }
 
-func (batch *BatchMutation) request(req *Req) {
-	counter := atomic.AddUint64(&batch.mutations, 1)
-RETRY:
-	_, err := batch.dc.Run(context.Background(), &req.gr)
-	if err != nil {
-		errString := err.Error()
-		// Irrecoverable
-		if strings.Contains(errString, "x509") || grpc.Code(err) == codes.Internal {
-			log.Fatal(errString)
-		}
-		fmt.Printf("Retrying req: %d. Error: %v\n", counter, errString)
-		time.Sleep(5 * time.Millisecond)
-		goto RETRY
-	}
-	req.reset()
+func (n *Node) ConnectTo(pred string, n1 Node) Edge {
+	e := Edge{}
+	e.nq.Subject = n.String()
+	e.nq.Predicate = pred
+	e.ConnectTo(n1)
+	return e
+}
+
+func (n *Node) Edge(pred string) Edge {
+	e := Edge{}
+	e.nq.Subject = n.String()
+	e.nq.Predicate = pred
+	return e
 }
 
-func (batch *BatchMutation) makeRequests() {
-	req := new(Req)
+type Edge struct {
+	nq protos.NQuad
+}
 
-	for n := range batch.nquads {
-		req.addMutation(n.nq, n.op)
-		if req.size() == batch.size {
-			batch.request(req)
-		}
-	}
+func NewEdge(nq protos.NQuad) Edge {
+	return Edge{nq}
+}
 
-	if req.size() > 0 {
-		batch.request(req)
+func (e *Edge) ConnectTo(n Node) error {
+	if e.nq.ObjectType > 0 {
+		return ErrValue
 	}
-	batch.wg.Done()
+	e.nq.ObjectId = n.String()
+	return nil
 }
 
-func (batch *BatchMutation) makeSchemaRequests() {
-	req := new(Req)
-LOOP:
-	for {
-		select {
-		case s, ok := <-batch.schema:
-			if !ok {
-				break LOOP
-			}
-			req.addSchema(s)
-		default:
-			start := time.Now()
-			if req.size() > 0 {
-				batch.request(req)
-			}
-			elapsedMillis := time.Since(start).Seconds() * 1e3
-			if elapsedMillis < 10 {
-				time.Sleep(time.Duration(int64(10-elapsedMillis)) * time.Millisecond)
-			}
+func validateStr(val string) error {
+	for idx, c := range val {
+		if c == '"' && (idx == 0 || val[idx-1] != '\\') {
+			return fmt.Errorf(`" must be preceded by a \ in object value`)
 		}
 	}
+	return nil
+}
 
-	if req.size() > 0 {
-		batch.request(req)
+func (e *Edge) SetValueString(val string) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
 	}
-	batch.wg.Done()
+	if err := validateStr(val); err != nil {
+		return err
+	}
+
+	v, err := types.ObjectValue(types.StringID, val)
+	if err != nil {
+		return err
+	}
+	e.nq.ObjectValue = v
+	e.nq.ObjectType = int32(types.StringID)
+	return nil
 }
 
-// NewBatchMutation is used to create a new batch.
-// size is the number of RDF's that are sent as part of one request to Dgraph.
-// pending is the number of concurrent requests to make to Dgraph server.
-func NewBatchMutation(ctx context.Context, client protos.DgraphClient,
-	size int, pending int) *BatchMutation {
-	bm := BatchMutation{
-		size:    size,
-		pending: pending,
-		nquads:  make(chan nquadOp, 2*size),
-		schema:  make(chan protos.SchemaUpdate, 2*size),
-		start:   time.Now(),
-		dc:      client,
-	}
-
-	for i := 0; i < pending; i++ {
-		bm.wg.Add(1)
-		go bm.makeRequests()
-	}
-	bm.wg.Add(1)
-	go bm.makeSchemaRequests()
-	return &bm
+func (e *Edge) SetValueInt(val int64) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
+	}
+	v, err := types.ObjectValue(types.IntID, val)
+	if err != nil {
+		return err
+	}
+	e.nq.ObjectValue = v
+	e.nq.ObjectType = int32(types.IntID)
+	return nil
 }
 
-// AddMutation is used to add a NQuad to a batch. It can either have SET or
-// DEL as Op(operation).
-func (batch *BatchMutation) AddMutation(nq protos.NQuad, op Op) error {
-	if err := checkNQuad(nq); err != nil {
+func (e *Edge) SetValueFloat(val float64) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
+	}
+	v, err := types.ObjectValue(types.FloatID, val)
+	if err != nil {
 		return err
 	}
-	batch.nquads <- nquadOp{nq: nq, op: op}
-	atomic.AddUint64(&batch.rdfs, 1)
+	e.nq.ObjectValue = v
+	e.nq.ObjectType = int32(types.FloatID)
 	return nil
 }
 
-// Flush waits for all pending requests to complete. It should always be called
-// after adding all the NQuads using batch.AddMutation().
-func (batch *BatchMutation) Flush() {
-	close(batch.nquads)
-	close(batch.schema)
-	batch.wg.Wait()
+func (e *Edge) SetValueBool(val bool) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
+	}
+	v, err := types.ObjectValue(types.BoolID, val)
+	if err != nil {
+		return err
+	}
+	e.nq.ObjectValue = v
+	e.nq.ObjectType = int32(types.BoolID)
+	return nil
 }
 
-// AddSchema is used to add a schema mutation.
-func (batch *BatchMutation) AddSchema(s protos.SchemaUpdate) error {
-	if err := checkSchema(s); err != nil {
+func (e *Edge) SetValuePassword(val string) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
+	}
+	v, err := types.ObjectValue(types.PasswordID, val)
+	if err != nil {
 		return err
 	}
-	batch.schema <- s
+	e.nq.ObjectValue = v
+	e.nq.ObjectType = int32(types.PasswordID)
 	return nil
 }
 
-// Counter keeps a track of various parameters about a batch mutation.
-type Counter struct {
-	// Number of RDF's processed by server.
-	Rdfs uint64
-	// Number of mutations processed by the server.
-	Mutations uint64
-	// Time elapsed sinze the batch started.
-	Elapsed time.Duration
+func (e *Edge) SetValueDatetime(dateTime time.Time) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
+	}
+	d, err := types.ObjectValue(types.DateTimeID, dateTime)
+	if err != nil {
+		return err
+	}
+	e.nq.ObjectValue = d
+	e.nq.ObjectType = int32(types.DateTimeID)
+	return nil
 }
 
-// Counter returns the current state of the BatchMutation.
-func (batch *BatchMutation) Counter() Counter {
-	return Counter{
-		Rdfs:      atomic.LoadUint64(&batch.rdfs),
-		Mutations: atomic.LoadUint64(&batch.mutations),
-		Elapsed:   time.Since(batch.start),
+func (e *Edge) SetValueGeoJson(json string) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
 	}
+	var g geom.T
+	// Parse the json
+	err := geojson.Unmarshal([]byte(json), &g)
+	if err != nil {
+		return err
+	}
+
+	geo, err := types.ObjectValue(types.GeoID, g)
+	if err != nil {
+		return err
+	}
+
+	e.nq.ObjectValue = geo
+	e.nq.ObjectType = int32(types.GeoID)
+	return nil
+}
+
+func (e *Edge) SetValueDefault(val string) error {
+	if len(e.nq.ObjectId) > 0 {
+		return ErrConnected
+	}
+	if err := validateStr(val); err != nil {
+		return err
+	}
+
+	v, err := types.ObjectValue(types.DefaultID, val)
+	if err != nil {
+		return err
+	}
+	e.nq.ObjectValue = v
+	e.nq.ObjectType = int32(types.StringID)
+	return nil
+}
+
+func (e *Edge) AddFacet(key, val string) {
+	e.nq.Facets = append(e.nq.Facets, &protos.Facet{
+		Key: key,
+		Val: val,
+	})
 }