Dgraph: K-Shortest Path with all predicates

main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"strings"

	"github.com/dgraph-io/dgo"
	"github.com/dgraph-io/dgo/protos/api"
	"github.com/pkg/errors"
	"google.golang.org/grpc"
)

var debug bool

func main() {
	flag.BoolVar(&debug, "v", false, "verbose")
	backend := flag.String("b", "localhost:9080", "Dgraph instance URL")
	director := flag.String("director", "Steven Spielberg", "the director")
	actor := flag.String("actor", "Jeff Goldblum", "the actor")
	flag.Parse()

	fmt.Printf("Finding the shortest path between director %s and actor %s\n", *director, *actor)

	conn, err := grpc.Dial(*backend, grpc.WithInsecure())
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()
	client := dgo.NewDgraphClient(api.NewDgraphClient(conn))

	ctx := context.Background()

	preds, err := predicates(ctx, client)
	if err != nil {
		log.Fatal(err)
	}

	steven, err := uidByNameAndPred(ctx, client, *director, "director.film")
	if err != nil {
		log.Fatal(err)
	}

	jeff, err := uidByNameAndPred(ctx, client, *actor, "actor.film")
	if err != nil {
		log.Fatal(err)
	}

	shortestPath(ctx, client, steven, jeff, preds)
}

func predicates(ctx context.Context, client *dgo.Dgraph) ([]string, error) {
	var data struct {
		Schema []struct{ Predicate, Type string }
	}
	if err := query(ctx, client, "schema{name\ntype}", &data); err != nil {
		return nil, err
	}

	var preds []string
	for _, pred := range data.Schema {
		if pred.Type == "password" || strings.HasPrefix(pred.Predicate, "dgraph.") {
			continue
		}
		preds = append(preds, pred.Predicate)
	}
	return preds, nil
}

func uidByNameAndPred(ctx context.Context, client *dgo.Dgraph, name, has string) (string, error) {
	const tmpl = `
	{
		req(func: eq(name@., %q)) @filter(has(%s)) {
			uid
		}
	}`

	var data struct{ Req []struct{ UID string } }
	if err := query(ctx, client, fmt.Sprintf(tmpl, name, has), &data); err != nil {
		return "", err
	}
	if len(data.Req) == 0 {
		return "", errors.Errorf("no matches were found")
	}
	return data.Req[0].UID, nil
}

func shortestPath(ctx context.Context, client *dgo.Dgraph, from, to string, preds []string) {
	const tmpl = `{
		path as shortest(from: %s, to: %s, depth: 2) {
		%s
		}
		path(func: uid(path)) {
			uid
			name@.
		}
	}`

	var b strings.Builder
	for _, p := range preds {
		fmt.Fprintf(&b, "\t\t<%s>\n", p)
	}

	var data struct {
		Path []struct {
			Name string `json:"name@."`
		}
	}
	if err := query(ctx, client, fmt.Sprintf(tmpl, from, to, b.String()), &data); err != nil {
		log.Fatal(err)
	}

	for _, p := range data.Path {
		if p.Name != "" {
			fmt.Println(p.Name)
		}
	}
}

func query(ctx context.Context, client *dgo.Dgraph, query string, dst interface{}) error {
	if debug {
		log.Printf("query: %s", query)
	}
	res, err := client.NewReadOnlyTxn().Query(ctx, query)
	if err != nil {
		return errors.Wrapf(err, "could not query Dgraph")
	}
	if debug {
		log.Printf("response: %s", res.GetJson())
	}

	if err := json.Unmarshal(res.GetJson(), &dst); err != nil {
		return errors.Wrap(err, "could not parse response")
	}
	return nil
}