FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Support bulk loader use-case to import unencrypted export and encrypt… · elasticjava/dgraph@d982be3 · GitHub

/ dgraph Public
forked from dgraph-io/dgraph

Commit d982be3

Browse files
parasssh
authored
Support bulk loader use-case to import unencrypted export and encrypt the result. (dgraph-io#5209)
Fixes DGRAPH-1254
1 parent ef2adf6 commit d982be3

3 files changed

Lines changed: 30 additions & 11 deletions

File tree

‎chunker/chunk.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -348,9 +348,10 @@ func slurpQuoted(r *bufio.Reader, out *bytes.Buffer) error {
348348
}
349349
}
350350

351-
// FileReader returns an open reader and file on the given file. Gzip-compressed input is detected
352-
// and decompressed automatically even without the gz extension. The caller is responsible for
353-
// calling the returned cleanup function when done with the reader.
351+
// FileReader returns an open reader on the given file. Gzip-compressed input is detected
352+
// and decompressed automatically even without the gz extension. The keyfile, if non-nil,
353+
// is used to decrypt the file. The caller is responsible for calling the returned cleanup
354+
// function when done with the reader.
354355
func FileReader(file string, keyfile string) (rd *bufio.Reader, cleanup func()) {
355356
var f *os.File
356357
var err error

‎dgraph/cmd/bulk/loader.go‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ type options struct {
6363
CustomTokenizers string
6464
NewUids bool
6565
ClientDir string
66+
Encrypted bool
6667

6768
MapShards int
6869
ReduceShards int
@@ -116,7 +117,7 @@ func newLoader(opt *options) *loader {
116117
readerChunkCh: make(chan *bytes.Buffer, opt.NumGoroutines),
117118
writeTs: getWriteTimestamp(zero),
118119
}
119-
st.schema = newSchemaStore(readSchema(opt.SchemaFile, opt.BadgerKeyFile), opt, st)
120+
st.schema = newSchemaStore(readSchema(opt), opt, st)
120121
ld := &loader{
121122
state: st,
122123
mappers: make([]*mapper, opt.NumGoroutines),
@@ -143,13 +144,18 @@ func getWriteTimestamp(zero *grpc.ClientConn) uint64 {
143144
}
144145
}
145146

146-
func readSchema(filename string, keyfile string) *schema.ParsedSchema {
147-
f, err := os.Open(filename)
147+
func readSchema(opt *options) *schema.ParsedSchema {
148+
f, err := os.Open(opt.SchemaFile)
148149
x.Check(err)
149150
defer f.Close()
151+
152+
keyfile := opt.BadgerKeyFile
153+
if !opt.Encrypted {
154+
keyfile = ""
155+
}
150156
r, err := enc.GetReader(keyfile, f)
151157
x.Check(err)
152-
if filepath.Ext(filename) == ".gz" {
158+
if filepath.Ext(opt.SchemaFile) == ".gz" {
153159
r, err = gzip.NewReader(r)
154160
x.Check(err)
155161
}
@@ -208,7 +214,11 @@ func (ld *loader) mapStage() {
208214
go func(file string) {
209215
defer thr.Done(nil)
210216

211-
r, cleanup := chunker.FileReader(file, ld.opt.BadgerKeyFile)
217+
keyfile := ld.opt.BadgerKeyFile
218+
if !ld.opt.Encrypted {
219+
keyfile = ""
220+
}
221+
r, cleanup := chunker.FileReader(file, keyfile)
212222
defer cleanup()
213223

214224
chunk := chunker.NewChunker(loadType, 1000)

‎dgraph/cmd/bulk/run.go‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@ func init() {
5858
"Location of schema file.")
5959
flag.String("format", "",
6060
"Specify file format (rdf or json) instead of getting it from filename.")
61+
flag.Bool("encrypted", false,
62+
"Flag to indicate whether schema and data files are encrypted.")
6163
flag.String("out", defaultOutDir,
6264
"Location to write the final dgraph data directories.")
6365
flag.Bool("replace_out", false,
@@ -101,9 +103,10 @@ func init() {
101103

102104
// Options around how to set up Badger.
103105
flag.String("encryption_key_file", "",
104-
"The file that stores the encryption key. The key size must be 16, 24, or 32 bytes long. "+
105-
"The key size determines the corresponding block size for AES encryption "+
106-
"(AES-128, AES-192, and AES-256 respectively). Enterprise feature.")
106+
"The file that stores the encryption key. The key size must be 16/24/32 bytes long."+
107+
" The key size indicates the chosen AES encryption (AES-128/192/256 respectively). "+
108+
" This key is used to encrypt the output data directories and to decrypt the input "+
109+
" schema and data files (if encrytped). Enterprise feature.")
107110
flag.Int("badger.compression_level", 1,
108111
"The compression level for Badger. A higher value uses more resources.")
109112
}
@@ -113,6 +116,7 @@ func run() {
113116
DataFiles: Bulk.Conf.GetString("files"),
114117
DataFormat: Bulk.Conf.GetString("format"),
115118
SchemaFile: Bulk.Conf.GetString("schema"),
119+
Encrypted: Bulk.Conf.GetBool("encrypted"),
116120
OutDir: Bulk.Conf.GetString("out"),
117121
ReplaceOutDir: Bulk.Conf.GetBool("replace_out"),
118122
TmpDir: Bulk.Conf.GetString("tmp"),
@@ -145,6 +149,10 @@ func run() {
145149
fmt.Printf("Cannot enable encryption: %s", x.ErrNotSupported)
146150
os.Exit(1)
147151
}
152+
if opt.Encrypted && opt.BadgerKeyFile == "" {
153+
fmt.Printf("Must use --encryption_key_file option with --encrypted option.\n")
154+
os.Exit(1)
155+
}
148156
if opt.SchemaFile == "" {
149157
fmt.Fprint(os.Stderr, "Schema file must be specified.\n")
150158
os.Exit(1)

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL