Golang DB.GetProto方法代码示例

// compareBiogoTree walks both a biogo tree and the range tree to determine if both
// contain the same values in the same order.
func compareBiogoTree(db *client.DB, biogoTree *llrb.Tree) error {
	rt := &proto.RangeTree{}
	if err := db.GetProto(keys.RangeTreeRoot, rt); err != nil {
		return err
	}
	return compareBiogoNode(db, biogoTree.Root, &rt.RootKey)
}

func countRangeReplicas(db *client.DB) (int, error) {
	desc := &roachpb.RangeDescriptor{}
	if err := db.GetProto(keys.RangeDescriptorKey(roachpb.KeyMin), desc); err != nil {
		return 0, err
	}
	return len(desc.Replicas), nil
}

// getConfig retrieves the configuration for the specified key. If the
// key is empty, all configurations are returned. Otherwise, the
// leading "/" path delimiter is stripped and the configuration
// matching the remainder is retrieved. Note that this will retrieve
// the default config if "key" is equal to "/", and will list all
// configs if "key" is equal to "". The body result contains a listing
// of keys and retrieval of a config. The output format is determined
// by the request header.
func getConfig(db *client.DB, configPrefix proto.Key, config gogoproto.Message,
	path string, r *http.Request) (body []byte, contentType string, err error) {
	// Scan all configs if the key is empty.
	if len(path) == 0 {
		var rows []client.KeyValue
		if rows, err = db.Scan(configPrefix, configPrefix.PrefixEnd(), maxGetResults); err != nil {
			return
		}
		if len(rows) == maxGetResults {
			log.Warningf("retrieved maximum number of results (%d); some may be missing", maxGetResults)
		}
		var prefixes []string
		for _, row := range rows {
			trimmed := bytes.TrimPrefix(row.Key, configPrefix)
			prefixes = append(prefixes, url.QueryEscape(string(trimmed)))
		}
		// Encode the response.
		body, contentType, err = util.MarshalResponse(r, prefixes, util.AllEncodings)
	} else {
		configkey := keys.MakeKey(configPrefix, proto.Key(path[1:]))
		if err = db.GetProto(configkey, config); err != nil {
			return
		}
		body, contentType, err = util.MarshalResponse(r, config, util.AllEncodings)
	}

	return
}

// compareBiogoNode compares a biogo node and a range tree node to determine if both
// contain the same values in the same order.  It recursively calls itself on
// both children if they exist.
func compareBiogoNode(db *client.DB, biogoNode *llrb.Node, key *proto.Key) error {
	// Retrieve the node form the range tree.
	rtNode := &proto.RangeTreeNode{}
	if err := db.GetProto(keys.RangeTreeNodeKey(*key), rtNode); err != nil {
		return err
	}

	bNode := &proto.RangeTreeNode{
		Key:       proto.Key(biogoNode.Elem.(Key)),
		ParentKey: proto.KeyMin,
		Black:     bool(biogoNode.Color),
	}
	if biogoNode.Left != nil {
		leftKey := proto.Key(biogoNode.Left.Elem.(Key))
		bNode.LeftKey = &leftKey
	}
	if biogoNode.Right != nil {
		rightKey := proto.Key(biogoNode.Right.Elem.(Key))
		bNode.RightKey = &rightKey
	}
	if err := nodesEqual(*key, *bNode, *rtNode); err != nil {
		return err
	}
	if rtNode.LeftKey != nil {
		if err := compareBiogoNode(db, biogoNode.Left, rtNode.LeftKey); err != nil {
			return err
		}
	}
	if rtNode.RightKey != nil {
		if err := compareBiogoNode(db, biogoNode.Right, rtNode.RightKey); err != nil {
			return err
		}
	}
	return nil
}

// GetTableDescriptor retrieves a table descriptor directly from the KV layer.
func GetTableDescriptor(kvDB *client.DB, database string, table string) *TableDescriptor {
	dbNameKey := MakeNameMetadataKey(keys.RootNamespaceID, database)
	gr, err := kvDB.Get(dbNameKey)
	if err != nil {
		panic(err)
	}
	if !gr.Exists() {
		panic("database missing")
	}
	dbDescID := ID(gr.ValueInt())

	tableNameKey := MakeNameMetadataKey(dbDescID, table)
	gr, err = kvDB.Get(tableNameKey)
	if err != nil {
		panic(err)
	}
	if !gr.Exists() {
		panic("table missing")
	}

	descKey := MakeDescMetadataKey(ID(gr.ValueInt()))
	desc := &Descriptor{}
	if err := kvDB.GetProto(descKey, desc); err != nil {
		panic("proto missing")
	}
	return desc.GetTable()
}

// getPermConfig fetches the permissions config for 'prefix'.
func getPermConfig(db *client.DB, prefix string) (*config.PermConfig, error) {
	config := &config.PermConfig{}
	if err := db.GetProto(keys.MakeKey(keys.ConfigPermissionPrefix, proto.Key(prefix)), config); err != nil {
		return nil, err
	}

	return config, nil
}

// loadTree loads the tree root and all of its nodes. It puts all of the nodes
// into a map.
func loadTree(t *testing.T, db *client.DB) (*roachpb.RangeTree, map[string]roachpb.RangeTreeNode) {
	tree := new(roachpb.RangeTree)
	if err := db.GetProto(keys.RangeTreeRoot, tree); err != nil {
		t.Fatal(err)
	}
	nodes := make(map[string]roachpb.RangeTreeNode)
	if tree.RootKey != nil {
		loadNodes(t, db, tree.RootKey, nodes)
	}
	return tree, nodes
}

// loadTree loads the tree root and all of its nodes. It puts all of the nodes
// into a map.
func loadTree(t *testing.T, db *client.DB) (storage.RangeTree, map[string]storage.RangeTreeNode) {
	var tree storage.RangeTree
	if err := db.GetProto(keys.RangeTreeRoot, &tree); err != nil {
		t.Fatal(err)
	}
	nodes := make(map[string]storage.RangeTreeNode)
	if tree.RootKey != nil {
		loadNodes(t, db, tree.RootKey, nodes)
	}
	return tree, nodes
}

// treesEqual compares the expectedTree and expectedNodes to the actual range
// tree stored in the db.
func treesEqual(db *client.DB, expected testRangeTree) error {
	// Compare the tree roots.
	actualTree := &proto.RangeTree{}
	if err := db.GetProto(keys.RangeTreeRoot, actualTree); err != nil {
		return err
	}
	if !reflect.DeepEqual(&expected.Tree, actualTree) {
		return util.Errorf("Range tree root is not as expected - expected:%+v - actual:%+v", expected.Tree, actualTree)
	}

	return treeNodesEqual(db, expected, expected.Tree.RootKey)
}

// setDefaultRangeMaxBytes sets the range-max-bytes value for the default zone.
func setDefaultRangeMaxBytes(t *testing.T, db *client.DB, maxBytes int64) {
	zone := &proto.ZoneConfig{}
	if err := db.GetProto(keys.ConfigZonePrefix, zone); err != nil {
		t.Fatal(err)
	}
	if zone.RangeMaxBytes == maxBytes {
		return
	}
	zone.RangeMaxBytes = maxBytes
	if err := db.Put(keys.ConfigZonePrefix, zone); err != nil {
		t.Fatal(err)
	}
}

// loadNodes fetches a node and recursively all of its children.
func loadNodes(t *testing.T, db *client.DB, key roachpb.RKey, nodes map[string]roachpb.RangeTreeNode) {
	node := new(roachpb.RangeTreeNode)
	if err := db.GetProto(keys.RangeTreeNodeKey(key), node); err != nil {
		t.Fatal(err)
	}
	nodes[node.Key.String()] = *node
	if node.LeftKey != nil {
		loadNodes(t, db, node.LeftKey, nodes)
	}
	if node.RightKey != nil {
		loadNodes(t, db, node.RightKey, nodes)
	}
}

// treeNodesEqual compares the expectedTree from the provided key to the actual
// nodes retrieved from the db.  It recursively calls itself on both left and
// right children if they exist.
func treeNodesEqual(db *client.DB, expected testRangeTree, key proto.Key) error {
	expectedNode, ok := expected.Nodes[string(key)]
	if !ok {
		return util.Errorf("Expected does not contain a node for %s", key)
	}
	actualNode := &proto.RangeTreeNode{}
	if err := db.GetProto(keys.RangeTreeNodeKey(key), actualNode); err != nil {
		return err
	}
	if err := nodesEqual(key, expectedNode, *actualNode); err != nil {
		return err
	}
	if expectedNode.LeftKey != nil {
		if err := treeNodesEqual(db, expected, *expectedNode.LeftKey); err != nil {
			return err
		}
	}
	if expectedNode.RightKey != nil {
		if err := treeNodesEqual(db, expected, *expectedNode.RightKey); err != nil {
			return err
		}
	}
	return nil
}

// Run a particular schema change and run some OLTP operations in parallel, as
// soon as the schema change starts executing its backfill.
func runSchemaChangeWithOperations(
	t *testing.T,
	sqlDB *gosql.DB,
	kvDB *client.DB,
	schemaChange string,
	maxValue int,
	keyMultiple int,
	descKey roachpb.Key,
	backfillNotification chan bool,
) {
	desc := &sqlbase.Descriptor{}
	if err := kvDB.GetProto(descKey, desc); err != nil {
		t.Fatal(err)
	}
	tableDesc := desc.GetTable()

	// Run the schema change in a separate goroutine.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		start := timeutil.Now()
		// Start schema change that eventually runs a backfill.
		if _, err := sqlDB.Exec(schemaChange); err != nil {
			t.Error(err)
		}
		t.Logf("schema change %s took %v", schemaChange, timeutil.Since(start))
		wg.Done()
	}()

	// Wait until the schema change backfill starts.
	<-backfillNotification

	// Run a variety of operations during the backfill.

	// Grabbing a schema change lease on the table will fail, disallowing
	// another schema change from being simultaneously executed.
	sc := csql.NewSchemaChangerForTesting(tableDesc.ID, 0, 0, *kvDB, nil)
	if l, err := sc.AcquireLease(); err == nil {
		t.Fatalf("schema change lease acquisition on table %d succeeded: %v", tableDesc.ID, l)
	}

	// Update some rows.
	var updatedKeys []int
	for i := 0; i < 10; i++ {
		k := rand.Intn(maxValue)
		v := maxValue + i + 1
		if _, err := sqlDB.Exec(`UPDATE t.test SET v = $2 WHERE k = $1`, k, v); err != nil {
			t.Fatal(err)
		}
		updatedKeys = append(updatedKeys, k)
	}

	// Reupdate updated values back to what they were before.
	for _, k := range updatedKeys {
		if _, err := sqlDB.Exec(`UPDATE t.test SET v = $2 WHERE k = $1`, k, maxValue-k); err != nil {
			t.Fatal(err)
		}
	}

	// Delete some rows.
	deleteStartKey := rand.Intn(maxValue - 10)
	for i := 0; i < 10; i++ {
		if _, err := sqlDB.Exec(`DELETE FROM t.test WHERE k = $1`, deleteStartKey+i); err != nil {
			t.Fatal(err)
		}
	}
	// Reinsert deleted rows.
	for i := 0; i < 10; i++ {
		k := deleteStartKey + i
		if _, err := sqlDB.Exec(`INSERT INTO t.test VALUES($1, $2)`, k, maxValue-k); err != nil {
			t.Fatal(err)
		}
	}

	// Insert some new rows.
	numInserts := 10
	for i := 0; i < numInserts; i++ {
		if _, err := sqlDB.Exec(`INSERT INTO t.test VALUES($1, $2)`, maxValue+i+1, maxValue+i+1); err != nil {
			t.Fatal(err)
		}
	}

	wg.Wait() // for schema change to complete.

	// Verify the number of keys left behind in the table to validate schema
	// change operations.
	tablePrefix := roachpb.Key(keys.MakeTablePrefix(uint32(tableDesc.ID)))
	tableEnd := tablePrefix.PrefixEnd()
	if kvs, err := kvDB.Scan(tablePrefix, tableEnd, 0); err != nil {
		t.Fatal(err)
	} else if e := keyMultiple * (maxValue + numInserts + 1); len(kvs) != e {
		t.Fatalf("expected %d key value pairs, but got %d", e, len(kvs))
	}

	// Delete the rows inserted.
	for i := 0; i < numInserts; i++ {
		if _, err := sqlDB.Exec(`DELETE FROM t.test WHERE k = $1`, maxValue+i+1); err != nil {
			t.Fatal(err)
//.........这里部分代码省略.........