Skip to content

Instantly share code, notes, and snippets.

@campoy
Last active Nov 9, 2019
Embed
What would you like to do?
Dgraph: K-Shortest Path with all predicates
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment