Merge branch 'work/srevin/clean-feature' into 'main'

feat: add clean support for shared dgraph instances

See merge request sorcero/community/dgbrx!1

README.md

@@ -26,10 +26,10 @@ Workflow 🔧
 -----------
 A general backup-restore workflow for dgraph is given below:
 ```bash
-dgbrx backup --url https://some-cool-url.region.gcp.cloud.dgraph.io/admin \
+dgbrx backup --url https://some-cool-url.region.gcp.cloud.dgraph.io \
   --api-key "SUPERSECRETAPIKEY"
 
-dgbrx restore --url https://another-cool-url.region.gcp.gcloud.dgraph.io/admin \
+dgbrx restore --url https://another-cool-url.region.gcp.gcloud.dgraph.io \
   --api-key "SUPERSECRETAPIKEYBUTDIFFERENTONE" \
   --json g01.json.gz \
   --schema g01.schema.gz

cmd/dgbrx/main.go

@@ -13,7 +13,7 @@ var logger = log.New(os.Stdout)
 var dgraphParams = []cli.Flag{
 	&cli.StringFlag{
 		Name:     "url",
-		Usage:    "Base URL to the dgraph instance, or the base URL to the admin dgraph instance",
+		Usage:    "Base URL to the dgraph instance, or the base URL to the dgraph instance",
 		EnvVars:  []string{"DGRAPH_API_URL"},
 		Required: true,
 	},

dgraph/admin.go

@@ -2,19 +2,23 @@ package dgraph
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/dgraph-io/dgo/v210"
+	"github.com/dgraph-io/dgo/v210/protos/api"
 	"io/ioutil"
 	"net/http"
+	"strings"
 	"time"
 )
 
-func doSlashRequest(dgraph *DGraph, query string) ([]byte, error) {
-	endpoint := fmt.Sprintf("%s/slash", dgraph.Endpoint)
+func doGraphRequest(dgraph *DGraph, query interface{}, suffix string) ([]byte, error) {
+	endpoint := fmt.Sprintf("%s/%s", dgraph.Endpoint, suffix)
 	logger.Debugf("Preparing request for backup of Dgraph Database: %s with query %s", endpoint, query)
 
-	marshal, err := json.Marshal(GenericQuery{Query: query})
+	marshal, err := json.Marshal(query)
 	if err != nil {
 		return nil, err
 	}
@@ -44,6 +48,14 @@ func doSlashRequest(dgraph *DGraph, query string) ([]byte, error) {
 	return d, nil
 }
 
+func doSlashRequest(dgraph *DGraph, query string) ([]byte, error) {
+	return doGraphRequest(dgraph, GenericQuery{Query: query}, "admin/slash")
+}
+
+func doQueryRequest(dgraph *DGraph, query string) ([]byte, error) {
+	return doGraphRequest(dgraph, map[string]string{"query": query}, "query")
+}
+
 func ExportBackupQueueJob(dgraph *DGraph) (*AdminExportPromise, error) {
 	logger.Debugf("Preparing backup queue")
 	d, err := doSlashRequest(dgraph, graphExportMutationQuery)
@@ -109,29 +121,77 @@ func AwaitExportBackup(dgraph *DGraph, e *AdminExportPromise) (*AdminExport, err
 
 }
 
-func DropAll(dgraph *DGraph) error {
-	d, err := json.Marshal(map[string]bool{"drop_all": true})
-	endpoint := fmt.Sprintf("%s/alter", dgraph.Endpoint)
+/*func DropAll(dgraph *DGraph) error {
+	logger.Debugf("Preparing drop all")
+
+
+	ctx := context.Background()
+	conn, err := dgo.DialCloud(dgraph.Endpoint, dgraph.ApiToken)
+	if err != nil { return err }
+	dg := dgo.NewDgraphClient(api.NewDgraphClient(conn))
+	err = dg.Login(ctx, "groot", "password")
+	if err != nil { return err }
+	err = dg.Alter(ctx, &api.Operation{DropAll: true})
 	if err != nil {
 		return err
 	}
-	req, err := http.NewRequest(
-		"POST",
-		endpoint,
-		bytes.NewBuffer(d),
-	)
-	req.Header.Set("Dg-Auth", dgraph.ApiToken)
-	req.Header.Set("Content-Type", "application/json")
+	logger.Infof("Drop databases completely.")
+	return nil
+}*/
 
+func DropPredicate(ctx context.Context, dg *dgo.Dgraph, predicate string) error {
+	if strings.HasPrefix(predicate, "dgraph") {
+		// some internal dgraph schema which we would like to omit
+		logger.Infof("Skipping predicate '%s'.", predicate)
+		return nil
+	}
+	logger.Debugf("Removing predicate '%s'", predicate)
+
+	err := dg.Alter(ctx, &api.Operation{DropAttr: predicate})
 	if err != nil {
 		return err
 	}
-	logger.Debugf("Authenticating")
+	logger.Infof("Removed predicate '%s'.", predicate)
+	return nil
+}
 
-	_, err = http.DefaultClient.Do(req)
+func DropAll(dgraph *DGraph) error {
+	logger.Debugf("Preparing to dropAll data")
+	d, err := doSlashRequest(dgraph, dropAllMutationQuery)
+	if err != nil {
+		return err
+	}
+	logger.Debugf("Fetching current schema")
+	d, err = doQueryRequest(dgraph, "schema {}")
+	if err != nil {
+		return err
+	}
+
+	admin := &PredicateQuery{}
+	logger.Debug(string(d))
+	err = json.Unmarshal(d, admin)
+	if err != nil {
+		return err
+	}
+	if len(admin.Error) != 0 {
+		for i := range admin.Error {
+			logger.Warn(admin.Error[i].Message)
+		}
+		return errors.New(admin.Error[0].Message)
+	}
+	logger.Debugf("Preparing to drop schema predicates")
+	ctx := context.Background()
+	conn, err := dgo.DialCloud(dgraph.Endpoint, dgraph.ApiToken)
 	if err != nil {
 		return err
 	}
-	logger.Debugf("Drop All request succeeded")
+	dg := dgo.NewDgraphClient(api.NewDgraphClient(conn))
+
+	for i := range admin.Data.Schema {
+		if err := DropPredicate(ctx, dg, admin.Data.Schema[i].Predicate); err != nil {
+			return err
+		}
+	}
+	logger.Infof("Drop all data completed successfully")
 	return nil
 }

dgraph/queries.go

@@ -19,3 +19,9 @@ const graphExportStatusCheckMutationQuery = `query {
     status
   }
 }`
+
+const dropAllMutationQuery = `mutation {
+  dropData(allData: true) {
+	response { code message }
+  }
+}`

dgraph/types.go

@@ -66,3 +66,16 @@ type AdminExport struct {
 type GenericQuery struct {
 	Query string `json:"query"`
 }
+
+type PredicateDefinition struct {
+	Predicate string `json:"predicate"`
+}
+
+type PredicateQueryPayload struct {
+	Schema []PredicateDefinition `json:"schema"`
+}
+
+type PredicateQuery struct {
+	Data  PredicateQueryPayload `json:"data"`
+	Error []GenericError        `json:"errors,omitempty"`
+}

go.mod

@@ -9,63 +9,30 @@ replace (
 )
 
 require (
-	github.com/chartmuseum/storage v0.11.0
+	github.com/baidubce/bce-sdk-go v0.9.86
 	github.com/codeclysm/extract/v3 v3.0.2
+	github.com/dgraph-io/dgo/v210 v210.0.0-20210825123656-d3f867fe9cc3
 	github.com/urfave/cli/v2 v2.3.0
 	github.com/withmandala/go-log v0.1.0
 )
 
 require (
-	cloud.google.com/go v0.93.3 // indirect
-	cloud.google.com/go/storage v1.16.1 // indirect
-	github.com/Azure/azure-sdk-for-go v57.1.0+incompatible // indirect
-	github.com/Azure/go-autorest v14.2.0+incompatible // indirect
-	github.com/Azure/go-autorest/autorest v0.11.20 // indirect
-	github.com/Azure/go-autorest/autorest/adal v0.9.13 // indirect
-	github.com/Azure/go-autorest/autorest/date v0.3.0 // indirect
-	github.com/Azure/go-autorest/logger v0.2.1 // indirect
-	github.com/Azure/go-autorest/tracing v0.6.0 // indirect
-	github.com/NetEase-Object-Storage/nos-golang-sdk v0.0.0-00010101000000-000000000000 // indirect
-	github.com/aliyun/aliyun-oss-go-sdk v2.1.10+incompatible // indirect
-	github.com/aws/aws-sdk-go v1.40.37 // indirect
-	github.com/baidubce/bce-sdk-go v0.9.86 // indirect
-	github.com/coreos/go-systemd v0.0.0-20180511133405-39ca1b05acc7 // indirect
-	github.com/coreos/go-systemd/v22 v22.0.0 // indirect
-	github.com/coreos/pkg v0.0.0-20160727233714-3ac0863d7acf // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
-	github.com/form3tech-oss/jwt-go v3.2.2+incompatible // indirect
-	github.com/gofrs/uuid v4.0.0+incompatible // indirect
-	github.com/gogo/protobuf v1.2.1 // indirect
-	github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/go-cmp v0.5.6 // indirect
-	github.com/google/go-querystring v1.0.0 // indirect
-	github.com/google/uuid v1.1.1 // indirect
-	github.com/googleapis/gax-go/v2 v2.1.0 // indirect
-	github.com/gophercloud/gophercloud v0.20.0 // indirect
 	github.com/h2non/filetype v1.0.6 // indirect
-	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/juju/errors v0.0.0-20181118221551-089d3ea4e4d5 // indirect
-	github.com/mozillazg/go-httpheader v0.2.1 // indirect
-	github.com/oracle/oci-go-sdk v24.3.0+incompatible // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/russross/blackfriday/v2 v2.0.1 // indirect
 	github.com/shurcooL/sanitized_anchor_name v1.0.0 // indirect
 	github.com/smartystreets/goconvey v1.6.6 // indirect
-	github.com/tencentyun/cos-go-sdk-v5 v0.7.30 // indirect
-	go.etcd.io/etcd v0.0.0-00010101000000-000000000000 // indirect
-	go.opencensus.io v0.23.0 // indirect
-	go.uber.org/atomic v1.3.2 // indirect
-	go.uber.org/multierr v1.1.0 // indirect
-	go.uber.org/zap v1.10.0 // indirect
+	github.com/stretchr/testify v1.7.0 // indirect
 	golang.org/x/crypto v0.0.0-20210817164053-32db794688a5 // indirect
 	golang.org/x/net v0.0.0-20210903162142-ad29c8ab022f // indirect
-	golang.org/x/oauth2 v0.0.0-20210819190943-2bc19b11175f // indirect
 	golang.org/x/sys v0.0.0-20210823070655-63515b42dcdf // indirect
 	golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 // indirect
 	golang.org/x/text v0.3.6 // indirect
-	golang.org/x/time v0.0.0-20191024005414-555d28b269f0 // indirect
-	google.golang.org/api v0.56.0 // indirect
-	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto v0.0.0-20210828152312-66f60bf46e71 // indirect
 	google.golang.org/grpc v1.40.0 // indirect
 	google.golang.org/protobuf v1.27.1 // indirect