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

     1  /*
     2   * Copyright 2019 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package testutil
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"math"
    23  
    24  	"github.com/dgraph-io/badger"
    25  	"github.com/dgraph-io/badger/options"
    26  	bpb "github.com/dgraph-io/badger/pb"
    27  	"github.com/dgraph-io/dgraph/posting"
    28  	"github.com/dgraph-io/dgraph/protos/pb"
    29  	"github.com/dgraph-io/dgraph/types"
    30  	"github.com/dgraph-io/dgraph/x"
    31  )
    32  
    33  // GetPValues reads the specified p directory and returns the values for the given
    34  // attribute in a map.
    35  // TODO(martinmr): See if this method can be simplified (e.g not use stream framework).
    36  func GetPValues(pdir, attr string, readTs uint64) (map[string]string, error) {
    37  	opt := badger.DefaultOptions(pdir).WithTableLoadingMode(options.MemoryMap).
    38  		WithReadOnly(true)
    39  	db, err := badger.OpenManaged(opt)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  	defer db.Close()
    44  
    45  	values := make(map[string]string)
    46  
    47  	stream := db.NewStreamAt(math.MaxUint64)
    48  	stream.ChooseKey = func(item *badger.Item) bool {
    49  		pk, err := x.Parse(item.Key())
    50  		x.Check(err)
    51  		switch {
    52  		case pk.Attr != attr:
    53  			return false
    54  		case pk.IsSchema():
    55  			return false
    56  		}
    57  		return pk.IsData()
    58  	}
    59  	stream.KeyToList = func(key []byte, it *badger.Iterator) (*bpb.KVList, error) {
    60  		pk, err := x.Parse(key)
    61  		x.Check(err)
    62  		pl, err := posting.ReadPostingList(key, it)
    63  		if err != nil {
    64  			return nil, err
    65  		}
    66  		var list bpb.KVList
    67  		err = pl.Iterate(readTs, 0, func(p *pb.Posting) error {
    68  			vID := types.TypeID(p.ValType)
    69  			src := types.ValueForType(vID)
    70  			src.Value = p.Value
    71  			str, err := types.Convert(src, types.StringID)
    72  			if err != nil {
    73  				fmt.Println(err)
    74  				return err
    75  			}
    76  			value := str.Value.(string)
    77  			list.Kv = append(list.Kv, &bpb.KV{
    78  				Key:   []byte(fmt.Sprintf("%#x", pk.Uid)),
    79  				Value: []byte(value),
    80  			})
    81  			return nil
    82  		})
    83  		return &list, err
    84  	}
    85  	stream.Send = func(list *bpb.KVList) error {
    86  		for _, kv := range list.Kv {
    87  			values[string(kv.Key)] = string(kv.Value)
    88  		}
    89  		return nil
    90  	}
    91  	if err := stream.Orchestrate(context.Background()); err != nil {
    92  		return nil, err
    93  	}
    94  	return values, err
    95  }