Support passing GraphQL schema to bulk loader.

dgraph/cmd/bulk/loader.go

@@ -27,6 +27,7 @@ import (
 	"log"
 	"os"
 	"path/filepath"
+	"strconv"
 	"sync"
 	"time"
 
@@ -47,6 +48,7 @@ type options struct {
 	DataFiles        string
 	DataFormat       string
 	SchemaFile       string
+	GqlSchemaFile    string
 	OutDir           string
 	ReplaceOutDir    bool
 	TmpDir           string
@@ -237,6 +239,9 @@ func (ld *loader) mapStage() {
 	}
 	x.Check(thr.Finish())
 
+	// Send the graphql triples
+	ld.processGqlSchema(loadType)
+
 	close(ld.readerChunkCh)
 	mapperWg.Wait()
 
@@ -251,6 +256,51 @@ func (ld *loader) mapStage() {
 	ld.xids = nil
 }
 
+func (ld *loader) processGqlSchema(loadType chunker.InputFormat) {
+	if ld.opt.GqlSchemaFile == "" {
+		return
+	}
+
+	f, err := os.Open(ld.opt.GqlSchemaFile)
+	x.Check(err)
+	defer f.Close()
+
+	key := ld.opt.EncryptionKey
+	if !ld.opt.Encrypted {
+		key = nil
+	}
+	r, err := enc.GetReader(key, f)
+	x.Check(err)
+	if filepath.Ext(ld.opt.GqlSchemaFile) == ".gz" {
+		r, err = gzip.NewReader(r)
+		x.Check(err)
+	}
+
+	buf, err := ioutil.ReadAll(r)
+	x.Check(err)
+
+	rdfSchema := `_:gqlschema <dgraph.type> "dgraph.graphql" .
+	_:gqlschema <dgraph.graphql.xid> "dgraph.graphql.schema" .
+	_:gqlschema <dgraph.graphql.schema> %s .
+	`
+
+	jsonSchema := `{
+		"dgraph.type": "dgraph.graphql",
+		"dgraph.graphql.xid": "dgraph.graphql.schema",
+		"dgraph.graphql.schema": %s
+	}`
+
+	gqlBuf := &bytes.Buffer{}
+	schema := strconv.Quote(string(buf))
+	switch loadType {
+	case chunker.RdfFormat:
+		x.Check2(gqlBuf.Write([]byte(fmt.Sprintf(rdfSchema, schema))))
+	case chunker.JsonFormat:
+		x.Check2(gqlBuf.Write([]byte(fmt.Sprintf(jsonSchema, schema))))
+	}
+	ld.readerChunkCh <- gqlBuf
+}
+
 func (ld *loader) reduceStage() {
 	ld.prog.setPhase(reducePhase)
 

dgraph/cmd/bulk/reduce.go

@@ -58,6 +58,7 @@ func (r *reducer) run() error {
 	x.AssertTrue(len(dirs) == r.opt.ReduceShards)
 	x.AssertTrue(len(r.opt.shardOutputDirs) == r.opt.ReduceShards)
 
+	r.dbs = make([]*badger.DB, r.opt.ReduceShards)
 	thr := y.NewThrottle(r.opt.NumReducers)
 	for i := 0; i < r.opt.ReduceShards; i++ {
 		if err := thr.Do(); err != nil {
@@ -128,7 +129,7 @@ func (r *reducer) createBadger(i int) *badger.DB {
 	// Zero out the key from memory.
 	opt.EncryptionKey = nil
 
-	r.dbs = append(r.dbs, db)
+	r.dbs[i] = db
 	return db
 }
 

dgraph/cmd/bulk/run.go

@@ -56,6 +56,7 @@ func init() {
 		"Location of *.rdf(.gz) or *.json(.gz) file(s) to load.")
 	flag.StringP("schema", "s", "",
 		"Location of schema file.")
+	flag.StringP("graphql_schema", "g", "", "Location of the GraphQL schema file.")
 	flag.String("format", "",
 		"Specify file format (rdf or json) instead of getting it from filename.")
 	flag.Bool("encrypted", false,
@@ -116,6 +117,7 @@ func run() {
 		DataFiles:              Bulk.Conf.GetString("files"),
 		DataFormat:             Bulk.Conf.GetString("format"),
 		SchemaFile:             Bulk.Conf.GetString("schema"),
+		GqlSchemaFile:          Bulk.Conf.GetString("graphql_schema"),
 		Encrypted:              Bulk.Conf.GetBool("encrypted"),
 		OutDir:                 Bulk.Conf.GetString("out"),
 		ReplaceOutDir:          Bulk.Conf.GetBool("replace_out"),