dgraph.go

package dgraph

import (
	"context"
	"errors"
	"strconv"
	"strings"
	"time"

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

type Host struct {
	Name       string `json:"hostName"`
	Server     string `json:"server"`
	Port       int    `json:"port"`
	IsDefault  bool   `json:"isDefault"`
	IsDisabled bool   `json:"IsDisabled"`

	// UserName  string
	// Password  string
}

type Instance struct {
	client *dgo.Dgraph
	host   Host
}

type DGraphDAO struct {
	HostName string
}

var (
	instances   map[string]*Instance
	defaultHost string
	configured  bool
)

// NewClient returns a new dgraph client for provided configuration.
func NewClient(h Host) (*dgo.Dgraph, error) {

	if strings.TrimSpace(h.Server) == "" {
		return nil, errors.New("host address can not be empty")
	}

	address := bindDgraphServerWithPort(h.Server, h.Port)

	// Dial a gRPC connection. The address to dial to can be configured when
	// setting up the dgraph cluster.
	dialOpts := []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithBlock(), // block till we connect to the server
		// grpc.WithTimeout(time.Second * 5),
		grpc.WithDefaultCallOptions(
			// grpc.UseCompressor(gzip.Name),
			grpc.WaitForReady(true),
		),
	}

	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*5) // wait for 5 seconds to connect to grpc server. Exit if the deadline exceeds.
	defer cancel()
	d, err := grpc.DialContext(ctx, address, dialOpts...)
	if err == context.DeadlineExceeded {
		return nil, errors.New("graphdb connect error, connection timed out for host " + address)
	}
	if err != nil {
		return nil, err
	}

	client := dgo.NewDgraphClient(api.NewDgraphClient(d))

	// Note: Supported in Enterprise version only
	// if h.UserName != "" {
	// 	if err = client.Login(context.TODO(), h.UserName, h.Password); err != nil {
	// 		return nil, err
	// 	}
	// }

	return client, nil
}

// NewInstance creates n new v2 instance of dgraph client. This instance can be saved in cache with host name as identifier for further operations.
func NewInstance(client *dgo.Dgraph, host Host) *Instance {
	return &Instance{
		client: client,
		host:   host,
	}
}

func InitInstances(configs []Host) error {

	if configured {
		return nil
	}

	instances = make(map[string]*Instance, len(configs))

	for _, host := range configs {
		if host.IsDisabled {
			continue
		}

		client, err := NewClient(host)
		if err != nil {
			return err
		}

		instances[host.Name] = &Instance{
			client: client,
			host:   host,
		}

		if host.IsDefault {
			defaultHost = host.Name
		}
	}

	configured = true

	return nil
}

// GetInstance returns a preconfigured dgraph instance from cache. If not present, returns an error.
func GetInstance(hostName string) (*Instance, error) {
	if hostName == "" {
		hostName = defaultHost
	}

	i, ok := instances[hostName]

	if !ok {
		return nil, errors.New("instance not found")
	}

	return i, nil
}

// GetDAO returns a dao instance to access and manipulate graph data and schema.
func GetDAO(hostName string) *DGraphDAO {
	return &DGraphDAO{HostName: hostName}
}

// CreateSchema sets the provided schema for the nodes data.
func (dg *DGraphDAO) CreateSchema(ctx context.Context, schema string) error {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return err
	}

	return instance.client.Alter(ctx, &api.Operation{Schema: schema})
}

// SetData sets the provided data as a node. Can be used to create or update a node.
//
// For update, the data must contain `uid` field.
//
// Set `commitNow` to true to commit or discard the changes immediately.
func (dg *DGraphDAO) SetData(ctx context.Context, data []byte, commitNow bool) error {
	return dg.mutate(ctx, data, &api.Mutation{
		SetJson:   data,
		CommitNow: commitNow,
	})
}

// DeleteData deletes the node or provided node attribute.
//
// Set `commitNow` to true to commit or discard the changes immediately.
func (dg *DGraphDAO) DeleteData(ctx context.Context, data []byte, commitNow bool) error {
	return dg.mutate(ctx, data, &api.Mutation{
		DeleteJson: data,
		CommitNow:  commitNow,
	})
}

// mutate creates or updates the node data.
func (dg *DGraphDAO) mutate(ctx context.Context, data []byte, mtn *api.Mutation) error {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return err
	}

	txn := instance.client.NewTxn()

	_, err = txn.Mutate(ctx, mtn)

	return err
}

// GetData returns the nodes matching to the provided query.
//
// query variables can be provided in `vars` param. Safe to provide `nil` if no variables required.
//
// The result is against the provided key in the query.
func (dg *DGraphDAO) GetData(ctx context.Context, query string, vars map[string]string) ([]byte, error) {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return nil, err
	}

	txn := instance.client.NewReadOnlyTxn().BestEffort()
	var res *api.Response
	if len(vars) == 0 {
		res, err = txn.Query(ctx, query)
	} else {
		res, err = txn.QueryWithVars(ctx, query, vars)
	}

	if err != nil {
		return nil, err
	}

	return res.GetJson(), nil
}

// DropSchema deletes the current schema along with the data.
func (dg *DGraphDAO) DropSchema(ctx context.Context) error {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return err
	}

	return instance.client.Alter(ctx, &api.Operation{DropOp: api.Operation_ALL})
}

// DropData deletes complete data but maintains the schema.
func (dg *DGraphDAO) DropData(ctx context.Context) error {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return err
	}

	return instance.client.Alter(ctx, &api.Operation{DropOp: api.Operation_DATA})
}

// DropAttr deletes a specific attribute completely from data and the schema.
func (dg *DGraphDAO) DropAttr(ctx context.Context, attr string) error {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return err
	}

	return instance.client.Alter(ctx, &api.Operation{DropOp: api.Operation_ATTR, DropValue: attr})
}

// DropEdge deletes the edges for the mentioned node. predicate is the name of relationship between the node.
//
// Ex. Persion1 `follows` Person2.
func (dg *DGraphDAO) DropEdge(ctx context.Context, uid string, predicates ...string) error {
	instance, err := GetInstance(dg.HostName)
	if err != nil {
		return err
	}

	mu := &api.Mutation{}
	dgo.DeleteEdges(mu, uid, predicates...)

	mu.CommitNow = true
	_, err = instance.client.NewTxn().Mutate(ctx, mu)
	return err
}

func bindDgraphServerWithPort(server string, port int) string {
	// if port is empty then use default port 9080(GRPC Port) & bind to server ip

	if port <= 0 || strings.TrimSpace(strconv.Itoa(port)) == "" {
		return server + ":9080"
	}

	return server + ":" + strconv.Itoa(port)
}