github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/ee/backup/restore.go (about)

     1  // +build !oss
     2  
     3  /*
     4   * Copyright 2019 Dgraph Labs, Inc. and Contributors
     5   *
     6   * Licensed under the Dgraph Community License (the "License"); you
     7   * may not use this file except in compliance with the License. You
     8   * may obtain a copy of the License at
     9   *
    10   *     https://github.com/dgraph-io/dgraph/blob/master/licenses/DCL.txt
    11   */
    12  
    13  package backup
    14  
    15  import (
    16  	"bufio"
    17  	"compress/gzip"
    18  	"encoding/binary"
    19  	"encoding/hex"
    20  	"fmt"
    21  	"io"
    22  	"math"
    23  	"path/filepath"
    24  
    25  	"github.com/dgraph-io/badger"
    26  	"github.com/dgraph-io/badger/options"
    27  	bpb "github.com/dgraph-io/badger/pb"
    28  	"github.com/pkg/errors"
    29  
    30  	"github.com/dgraph-io/dgraph/posting"
    31  	"github.com/dgraph-io/dgraph/protos/pb"
    32  	"github.com/dgraph-io/dgraph/x"
    33  )
    34  
    35  // RunRestore calls badger.Load and tries to load data into a new DB.
    36  func RunRestore(pdir, location, backupId string) (uint64, error) {
    37  	// Scan location for backup files and load them. Each file represents a node group,
    38  	// and we create a new p dir for each.
    39  	return Load(location, backupId, func(r io.Reader, groupId int, preds predicateSet) error {
    40  		dir := filepath.Join(pdir, fmt.Sprintf("p%d", groupId))
    41  		db, err := badger.OpenManaged(badger.DefaultOptions(dir).
    42  			WithSyncWrites(false).
    43  			WithTableLoadingMode(options.MemoryMap).
    44  			WithValueThreshold(1 << 10).
    45  			WithNumVersionsToKeep(math.MaxInt32))
    46  		if err != nil {
    47  			return err
    48  		}
    49  		defer db.Close()
    50  		fmt.Printf("Restoring groupId: %d\n", groupId)
    51  		if !pathExist(dir) {
    52  			fmt.Println("Creating new db:", dir)
    53  		}
    54  		gzReader, err := gzip.NewReader(r)
    55  		if err != nil {
    56  			return nil
    57  		}
    58  		return loadFromBackup(db, gzReader, preds)
    59  	})
    60  }
    61  
    62  // loadFromBackup reads the backup, converts the keys and values to the required format,
    63  // and loads them to the given badger DB.
    64  func loadFromBackup(db *badger.DB, r io.Reader, preds predicateSet) error {
    65  	br := bufio.NewReaderSize(r, 16<<10)
    66  	unmarshalBuf := make([]byte, 1<<10)
    67  
    68  	loader := db.NewKVLoader(16)
    69  	for {
    70  		var sz uint64
    71  		err := binary.Read(br, binary.LittleEndian, &sz)
    72  		if err == io.EOF {
    73  			break
    74  		} else if err != nil {
    75  			return err
    76  		}
    77  
    78  		if cap(unmarshalBuf) < int(sz) {
    79  			unmarshalBuf = make([]byte, sz)
    80  		}
    81  
    82  		if _, err = io.ReadFull(br, unmarshalBuf[:sz]); err != nil {
    83  			return err
    84  		}
    85  
    86  		list := &bpb.KVList{}
    87  		if err := list.Unmarshal(unmarshalBuf[:sz]); err != nil {
    88  			return err
    89  		}
    90  
    91  		for _, kv := range list.Kv {
    92  			if len(kv.GetUserMeta()) != 1 {
    93  				return errors.Errorf(
    94  					"Unexpected meta: %v for key: %s", kv.UserMeta, hex.Dump(kv.Key))
    95  			}
    96  
    97  			restoreKey, err := fromBackupKey(kv.Key)
    98  			if err != nil {
    99  				return err
   100  			}
   101  
   102  			// Filter keys using the preds set. Do not do this filtering for type keys
   103  			// as they are meant to be in every group and their Attr value does not
   104  			// match a predicate name.
   105  			parsedKey, err := x.Parse(restoreKey)
   106  			if err != nil {
   107  				return errors.Wrapf(err, "could not parse key %s", hex.Dump(restoreKey))
   108  			}
   109  			if _, ok := preds[parsedKey.Attr]; !parsedKey.IsType() && !ok {
   110  				continue
   111  			}
   112  
   113  			var restoreVal []byte
   114  			switch kv.GetUserMeta()[0] {
   115  			case posting.BitEmptyPosting, posting.BitCompletePosting, posting.BitDeltaPosting:
   116  				var err error
   117  				backupPl := &pb.BackupPostingList{}
   118  				if err := backupPl.Unmarshal(kv.Value); err != nil {
   119  					return errors.Wrapf(err, "while reading backup posting list")
   120  				}
   121  				restoreVal, err = posting.FromBackupPostingList(backupPl).Marshal()
   122  				if err != nil {
   123  					return errors.Wrapf(err, "while converting backup posting list")
   124  				}
   125  
   126  			case posting.BitSchemaPosting:
   127  				restoreVal = kv.Value
   128  
   129  			default:
   130  				return errors.Errorf(
   131  					"Unexpected meta %d for key %s", kv.UserMeta[0], hex.Dump(kv.Key))
   132  			}
   133  
   134  			kv.Key = restoreKey
   135  			kv.Value = restoreVal
   136  			if err := loader.Set(kv); err != nil {
   137  				return err
   138  			}
   139  		}
   140  	}
   141  
   142  	if err := loader.Finish(); err != nil {
   143  		return err
   144  	}
   145  
   146  	return nil
   147  }
   148  
   149  func fromBackupKey(key []byte) ([]byte, error) {
   150  	backupKey := &pb.BackupKey{}
   151  	if err := backupKey.Unmarshal(key); err != nil {
   152  		return nil, errors.Wrapf(err, "while reading backup key %s", hex.Dump(key))
   153  	}
   154  	return x.FromBackupKey(backupKey), nil
   155  }