Tool to debug posting lists in RocksDB

tools/dlist/.gitignore

+/dlist

tools/dlist/main.go

-package dlist
+package main
 
 import (
+	"bytes"
 	"flag"
+	"fmt"
+	"strconv"
 
+	"github.com/Sirupsen/logrus"
 	"github.com/dgraph-io/dgraph/posting"
-	"github.com/dgraph-io/dgraph/store"
+	"github.com/dgraph-io/dgraph/posting/types"
+	"github.com/dgraph-io/dgraph/store/rocksdb"
+	"github.com/dgraph-io/dgraph/uid"
 	"github.com/dgraph-io/dgraph/x"
 )
 
@@ -12,23 +18,88 @@ var glog = x.Log("dlist")
 
 var dir = flag.String("dir", "", "Directory containing ")
 var xid = flag.String("xid", "", "Get posting list for xid")
-var uid = flag.String("uid", "", "Get posting list for uid")
+var suid = flag.String("uid", "", "Get posting list for uid")
 var attr = flag.String("attr", "", "Get posting list for attribute")
+var count = flag.Bool("count", false, "Only output number of results."+
+	" Useful for range scanning with attribute.")
+
+func output(val []byte) {
+	pl := types.GetRootAsPostingList(val, 0)
+	fmt.Printf("Found posting list of length: %v\n", pl.PostingsLength())
+	var p types.Posting
+	for i := 0; i < pl.PostingsLength(); i++ {
+		if !pl.Postings(&p, i) {
+			glog.WithField("i", i).Fatal("Unable to get posting")
+		}
+		fmt.Printf("[%v] Uid: [%#x] Value: [%s]\n",
+			i, p.Uid(), string(p.ValueBytes()))
+	}
+}
+
+func scanOverAttr(db *rocksdb.DB) {
+	ro := rocksdb.NewReadOptions()
+	ro.SetFillCache(false)
+
+	prefix := []byte(*attr)
+	itr := db.NewIterator(ro)
+	itr.Seek(prefix)
+
+	num := 0
+	for itr = itr; itr.Valid(); itr.Next() {
+		if !bytes.HasPrefix(itr.Key(), prefix) {
+			break
+		}
+		if !*count {
+			fmt.Printf("\nkey: %#x\n", itr.Key())
+			output(itr.Value())
+		}
+		num += 1
+	}
+	if err := itr.GetError(); err != nil {
+		glog.WithError(err).Fatal("While iterating")
+	}
+	fmt.Printf("Number of keys found: %v\n", num)
+}
 
 func main() {
+	logrus.SetLevel(logrus.ErrorLevel)
+
 	flag.Parse()
-	var s store.Store
-	s.Init(*dir)
-	defer s.Close()
+	opt := rocksdb.NewOptions()
+	db, err := rocksdb.Open(*dir, opt)
+	defer db.Close()
 
 	var key []byte
-	if len(*uid) > 0 && len(*attr) > 0 {
-		key = posting.Key(*uid, *attr)
+	if len(*suid) > 0 && len(*attr) > 0 {
+		u, rerr := strconv.ParseUint(*suid, 0, 64)
+		if rerr != nil {
+			glog.WithError(rerr).Fatal("While parsing uid")
+		}
+		key = posting.Key(u, *attr)
 
 	} else if len(*attr) > 0 {
-		glog.Fatal("Not handling this yet.")
+		scanOverAttr(db)
+		return
+
+	} else if len(*suid) > 0 {
+		u, rerr := strconv.ParseUint(*suid, 0, 64)
+		if rerr != nil {
+			glog.WithError(rerr).Fatal("While parsing uid")
+		}
+		key = posting.Key(u, "_xid_")
+
+	} else if len(*xid) > 0 {
+		key = uid.StringKey(*xid)
+
+	} else {
+		glog.Fatal("Invalid request.")
+	}
+	fmt.Printf("key: %#x\n", key)
 
-	} else if len(*uid) > 0 {
-		key = posting.Key(*uid, "_xid_")
+	ropt := rocksdb.NewReadOptions()
+	val, err := db.Get(ropt, key)
+	if err != nil {
+		glog.WithError(err).Fatal("Unable to get key")
 	}
+	output(val)
 }

uid/assigner.go

@@ -187,12 +187,12 @@ func assignNew(pl *posting.List, xid string, instanceIdx uint64,
 	return uid, rerr
 }
 
-func stringKey(xid string) []byte {
+func StringKey(xid string) []byte {
 	return []byte("_uid_|" + xid)
 }
 
 func Get(xid string) (uid uint64, rerr error) {
-	key := stringKey(xid)
+	key := StringKey(xid)
 	pl := posting.GetOrCreate(key, uidStore)
 	if pl.Length() == 0 {
 		return 0, fmt.Errorf("xid: %v doesn't have any uid assigned.", xid)
@@ -210,7 +210,7 @@ func Get(xid string) (uid uint64, rerr error) {
 func GetOrAssign(xid string, instanceIdx uint64,
 	numInstances uint64) (uid uint64, rerr error) {
 
-	key := stringKey(xid)
+	key := StringKey(xid)
 	pl := posting.GetOrCreate(key, uidStore)
 	if pl.Length() == 0 {
 		return assignNew(pl, xid, instanceIdx, numInstances)