dgraph.go

package dgraph

import (
	"errors"
	"fmt"
	"reflect"
	"strings"
	"sync/atomic"

	"github.com/mitchellh/reflectwalk"
)

const (
	dgraphTypePredicate = "dgraph.type"
	dgraphTagName       = "dgraph"
)

// getPredicate gets the dgraph predicate name from the JSON struct tag.
func getPredicate(field *reflect.StructField) string {
	jsonTags := strings.Split(field.Tag.Get("json"), ",")
	return jsonTags[0]
}

// SetTypes recursively walks all structures in data and sets the value of the
// `dgraph.type` struct field. The type, in order of preference, is either the
// value of the `dgraph` struct tag on the `dgraph.type` struct field, or the
// struct name.
func SetTypes(data interface{}) error {
	w := typeWalker{}
	return reflectwalk.Walk(data, w)
}

type typeWalker struct{}

func (w typeWalker) Struct(v reflect.Value) error {
	vType := v.Type()
	nodeType := vType.Name()

	for i := 0; i < v.NumField(); i++ {
		field := vType.Field(i)
		fieldVal := v.Field(i)

		if getPredicate(&field) == dgraphTypePredicate {
			if !fieldVal.CanSet() {
				return fmt.Errorf("dgraph.type not settable on %s.%s", nodeType, field.Name) // did you pass pointer?
			}

			dgraphTag := field.Tag.Get(dgraphTagName)
			if dgraphTag != "" {
				nodeType = dgraphTag
			}
			switch field.Type.Kind() {
			case reflect.String:
				fieldVal.SetString(nodeType)
			case reflect.Slice:
				if field.Type.Elem().Kind() != reflect.String {
					return errors.New(`"dgraph.type" field is not a slice of strings`)
				}
				fieldVal.Set(reflect.ValueOf([]string{nodeType}))
			default:
				return errors.New(`unsupported type for "dgraph.type" predicate`)
			}

			break
		}
	}
	return nil
}

func (w typeWalker) StructField(f reflect.StructField, v reflect.Value) error {
	return nil
}

// SetUids recursively walks all structures in data and sets the value of the
// `uid` struct field based on the uids map. A map of Uids is returned in Dgraph
// mutate calls.
func SetUids(data interface{}, uids map[string]string) error {
	w := setUidWalker{uids: uids}
	return reflectwalk.Walk(data, w)
}

type setUidWalker struct {
	uids map[string]string
}

func (w setUidWalker) Struct(v reflect.Value) error {
	return nil
}

func (w setUidWalker) StructField(f reflect.StructField, v reflect.Value) error {
	if v.Kind() != reflect.String {
		return nil
	}

	predicate := getPredicate(&f)
	blankUid := v.String()

	if predicate == "uid" && strings.HasPrefix(blankUid, "_:") {
		uid, ok := w.uids[blankUid[2:]]
		if ok && v.CanSet() {
			v.Set(reflect.ValueOf(uid))
		}
		if !v.CanSet() {
			return fmt.Errorf("cannot set %s/%s", predicate, blankUid)
		}
	}
	return nil
}

// overflow is OK
var blankuid int32 = 0

func blankUid() string {
	i := atomic.AddInt32(&blankuid, 1)
	return fmt.Sprintf("_:%d", i)
}

// GenUids recursively walks all structures in data and sets the value of the
// `uid` struct field to a valid random blank uid. Only nodes with a blank uid
// set before mutate will be able to associate the blank uid with the persisted
// Uid after mutate. To Swap the random blank uids with valid uids, call SetUids
// on data after the mutation.
func GenUids(data interface{}) error {
	w := genUidWalker{}
	return reflectwalk.Walk(data, w)
}

type genUidWalker struct {
	uids map[string]string
}

func (w genUidWalker) Struct(v reflect.Value) error {
	return nil
}

func (w genUidWalker) StructField(f reflect.StructField, v reflect.Value) error {
	if v.Kind() != reflect.String {
		return nil
	}

	predicate := getPredicate(&f)
	uid := v.String()

	if predicate == "uid" && uid == "" {
		if !v.CanSet() {
			return fmt.Errorf("cannot set uid")
		}
		v.Set(reflect.ValueOf(blankUid()))
	}
	return nil
}