Golang client.Batch類代碼示例

// Future home of the asynchronous schema changer that picks up
// queued schema changes and processes them.
//
// applyMutations applies the queued mutations for a table.
func (p *planner) applyMutations(tableDesc *TableDescriptor) error {
	if len(tableDesc.Mutations) == 0 {
		return nil
	}
	newTableDesc := proto.Clone(tableDesc).(*TableDescriptor)
	p.applyUpVersion(newTableDesc)
	// Make all mutations active.
	for _, mutation := range newTableDesc.Mutations {
		newTableDesc.makeMutationComplete(mutation)
	}
	newTableDesc.Mutations = nil
	if err := newTableDesc.Validate(); err != nil {
		return err
	}

	b := client.Batch{}
	if err := p.backfillBatch(&b, tableDesc, newTableDesc); err != nil {
		return err
	}

	b.Put(MakeDescMetadataKey(newTableDesc.GetID()), wrapDescriptor(newTableDesc))

	if err := p.txn.Run(&b); err != nil {
		return convertBatchError(newTableDesc, b, err)
	}
	p.notifyCompletedSchemaChange(newTableDesc.ID)
	return nil
}

// writeDescriptor takes a Table or Database descriptor and writes it
// if needed, incrementing the descriptor counter.
func (p *planner) writeDescriptor(plainKey descriptorKey, descriptor descriptorProto, ifNotExists bool) error {
	key := plainKey.Key()
	// Check whether key exists.
	gr, err := p.txn.Get(key)
	if err != nil {
		return err
	}

	if gr.Exists() {
		if ifNotExists {
			// Noop.
			return nil
		}
		// Key exists, but we don't want it to: error out.
		return fmt.Errorf("%s %q already exists", descriptor.TypeName(), plainKey.Name())
	}

	// Increment unique descriptor counter.
	if ir, err := p.txn.Inc(keys.DescIDGenerator, 1); err == nil {
		descriptor.SetID(structured.ID(ir.ValueInt() - 1))
	} else {
		return err
	}

	// TODO(pmattis): The error currently returned below is likely going to be
	// difficult to interpret.
	// TODO(pmattis): Need to handle if-not-exists here as well.
	descKey := structured.MakeDescMetadataKey(descriptor.GetID())
	b := client.Batch{}
	b.CPut(key, descKey, nil)
	b.CPut(descKey, descriptor, nil)
	return p.txn.Run(&b)
}

func runPut(cmd *cobra.Command, args []string) {
	if len(args) == 0 || len(args)%2 == 1 {
		cmd.Usage()
		return
	}

	count := len(args) / 2

	keys := make([]string, 0, count)
	values := make([]string, 0, count)
	for i := 0; i < len(args); i += 2 {
		keys = append(keys, unquoteArg(args[i], true /* disallow system keys */))
		values = append(values, unquoteArg(args[i+1], false))
	}

	kvDB := makeDBClient()
	if kvDB == nil {
		return
	}
	var b client.Batch
	for i := 0; i < count; i++ {
		b.Put(keys[i], values[i])
	}
	if err := kvDB.Run(&b); err != nil {
		fmt.Fprintf(osStderr, "put failed: %s\n", err)
		osExit(1)
		return
	}
}

// Truncate deletes all rows from a table.
// Privileges: WRITE on table.
//   Notes: postgres requires TRUNCATE.
//          mysql requires DROP (for mysql >= 5.1.16, DELETE before that).
func (p *planner) Truncate(n *parser.Truncate) (planNode, error) {
	b := client.Batch{}

	for _, tableQualifiedName := range n.Tables {
		tableDesc, err := p.getTableDesc(tableQualifiedName)
		if err != nil {
			return nil, err
		}

		if err := p.checkPrivilege(tableDesc, privilege.WRITE); err != nil {
			return nil, err
		}

		tablePrefix := MakeTablePrefix(tableDesc.ID)

		// Delete rows and indexes starting with the table's prefix.
		tableStartKey := proto.Key(tablePrefix)
		tableEndKey := tableStartKey.PrefixEnd()
		if log.V(2) {
			log.Infof("DelRange %q - %q", tableStartKey, tableEndKey)
		}
		b.DelRange(tableStartKey, tableEndKey)
	}

	if err := p.txn.Run(&b); err != nil {
		return nil, err
	}

	// TODO(tamird/pmattis): return the number of affected rows
	return &valuesNode{}, nil
}

// deleteRow adds to the batch the kv operations necessary to delete a table row
// with the given values.
func (rd *rowDeleter) deleteRow(b *client.Batch, values []parser.Datum) error {
	if err := rd.fks.checkAll(values); err != nil {
		return err
	}

	primaryIndexKey, secondaryIndexEntries, err := rd.helper.encodeIndexes(rd.fetchColIDtoRowIndex, values)
	if err != nil {
		return err
	}

	for _, secondaryIndexEntry := range secondaryIndexEntries {
		if log.V(2) {
			log.Infof("Del %s", secondaryIndexEntry.Key)
		}
		b.Del(secondaryIndexEntry.Key)
	}

	// Delete the row.
	rd.startKey = roachpb.Key(primaryIndexKey)
	rd.endKey = rd.startKey.PrefixEnd()
	if log.V(2) {
		log.Infof("DelRange %s - %s", rd.startKey, rd.endKey)
	}
	b.DelRange(&rd.startKey, &rd.endKey, false)
	rd.startKey, rd.endKey = nil, nil

	return nil
}

// Truncate deletes all rows from a table.
// Privileges: WRITE on table.
//   Notes: postgres requires TRUNCATE.
//          mysql requires DROP (for mysql >= 5.1.16, DELETE before that).
func (p *planner) Truncate(n *parser.Truncate) (planNode, error) {
	b := client.Batch{}

	for _, tableQualifiedName := range n.Tables {
		tableDesc, err := p.getTableDesc(tableQualifiedName)
		if err != nil {
			return nil, err
		}

		if !tableDesc.HasPrivilege(p.user, parser.PrivilegeWrite) {
			return nil, fmt.Errorf("user %s does not have %s privilege on table %s",
				p.user, parser.PrivilegeWrite, tableDesc.Name)
		}

		tablePrefix := encodeTablePrefix(tableDesc.ID)

		// Delete rows and indexes starting with the table's prefix.
		tableStartKey := proto.Key(tablePrefix)
		tableEndKey := tableStartKey.PrefixEnd()
		if log.V(2) {
			log.Infof("DelRange %q - %q", tableStartKey, tableEndKey)
		}
		b.DelRange(tableStartKey, tableEndKey)
	}

	if err := p.db.Run(&b); err != nil {
		return nil, err
	}

	// TODO(tamird/pmattis): return the number of affected rows
	return &valuesNode{}, nil
}

// Future home of the aysynchronous schema changer that picks up
// queued schema changes and processes them.
//
// applyMutations applies the queued mutations for a table.
// TODO(vivek): Eliminate the need to pass in tableName.
func (p *planner) applyMutations(tableDesc *TableDescriptor, tableName *parser.QualifiedName) error {
	if len(tableDesc.Mutations) == 0 {
		return nil
	}
	newTableDesc := proto.Clone(tableDesc).(*TableDescriptor)
	// Make all mutations active.
	for _, mutation := range newTableDesc.Mutations {
		newTableDesc.makeMutationComplete(mutation)
	}
	newTableDesc.Mutations = nil
	if err := newTableDesc.Validate(); err != nil {
		return err
	}

	b := client.Batch{}
	if err := p.backfillBatch(&b, tableName, tableDesc, newTableDesc); err != nil {
		return err
	}
	// TODO(pmattis): This is a hack. Remove when schema change operations work
	// properly.
	p.hackNoteSchemaChange(newTableDesc)

	b.Put(MakeDescMetadataKey(newTableDesc.GetID()), wrapDescriptor(newTableDesc))

	if err := p.txn.Run(&b); err != nil {
		return convertBatchError(newTableDesc, b, err)
	}
	return nil
}

// StoreData writes the supplied time series data to the cockroach server.
// Stored data will be sampled at the supplied resolution.
func (db *DB) StoreData(r Resolution, data []tspb.TimeSeriesData) error {
	var kvs []roachpb.KeyValue

	// Process data collection: data is converted to internal format, and a key
	// is generated for each internal message.
	for _, d := range data {
		idatas, err := d.ToInternal(r.KeyDuration(), r.SampleDuration())
		if err != nil {
			return err
		}
		for _, idata := range idatas {
			var value roachpb.Value
			if err := value.SetProto(&idata); err != nil {
				return err
			}
			kvs = append(kvs, roachpb.KeyValue{
				Key:   MakeDataKey(d.Name, d.Source, r, idata.StartTimestampNanos),
				Value: value,
			})
		}
	}

	// Send the individual internal merge requests.
	b := client.Batch{}
	for _, kv := range kvs {
		b.AddRawRequest(&roachpb.MergeRequest{
			Span: roachpb.Span{
				Key: kv.Key,
			},
			Value: kv.Value,
		})
	}

	return db.db.Run(&b)
}

// Truncate deletes all rows from a table.
// Privileges: DROP on table.
//   Notes: postgres requires TRUNCATE.
//          mysql requires DROP (for mysql >= 5.1.16, DELETE before that).
func (p *planner) Truncate(n *parser.Truncate) (planNode, error) {
	b := client.Batch{}
	for _, tableQualifiedName := range n.Tables {
		tableDesc, err := p.getTableDesc(tableQualifiedName)
		if err != nil {
			return nil, err
		}

		if err := p.checkPrivilege(tableDesc, privilege.DROP); err != nil {
			return nil, err
		}

		tablePrefix := keys.MakeTablePrefix(uint32(tableDesc.ID))

		// Delete rows and indexes starting with the table's prefix.
		tableStartKey := proto.Key(tablePrefix)
		tableEndKey := tableStartKey.PrefixEnd()
		if log.V(2) {
			log.Infof("DelRange %s - %s", prettyKey(tableStartKey, 0), prettyKey(tableEndKey, 0))
		}
		b.DelRange(tableStartKey, tableEndKey)
	}

	if err := p.txn.Run(&b); err != nil {
		return nil, err
	}

	return &valuesNode{}, nil
}

// Truncate deletes all rows from a table.
// Privileges: DROP on table.
//   Notes: postgres requires TRUNCATE.
//          mysql requires DROP (for mysql >= 5.1.16, DELETE before that).
func (p *planner) Truncate(n *parser.Truncate) (planNode, *roachpb.Error) {
	b := client.Batch{}
	for _, tableQualifiedName := range n.Tables {
		tableDesc, pErr := p.getTableLease(tableQualifiedName)
		if pErr != nil {
			return nil, pErr
		}

		if err := p.checkPrivilege(&tableDesc, privilege.DROP); err != nil {
			return nil, roachpb.NewError(err)
		}

		tablePrefix := keys.MakeTablePrefix(uint32(tableDesc.ID))

		// Delete rows and indexes starting with the table's prefix.
		tableStartKey := roachpb.Key(tablePrefix)
		tableEndKey := tableStartKey.PrefixEnd()
		if log.V(2) {
			log.Infof("DelRange %s - %s", tableStartKey, tableEndKey)
		}
		b.DelRange(tableStartKey, tableEndKey, false)
	}

	if pErr := p.txn.Run(&b); pErr != nil {
		return nil, pErr
	}

	return &emptyNode{}, nil
}

func runDel(cmd *cobra.Command, args []string) {
	if len(args) == 0 {
		cmd.Usage()
		return
	}

	count := len(args)

	keys := make([]string, 0, count)
	for i := 0; i < count; i++ {
		keys = append(keys, unquoteArg(args[i], true /* disallow system keys */))
	}

	kvDB := makeDBClient()
	if kvDB == nil {
		return
	}
	var b client.Batch
	for i := 0; i < count; i++ {
		b.Del(keys[i])
	}
	if err := kvDB.Run(&b); err != nil {
		fmt.Fprintf(osStderr, "delete failed: %s\n", err)
		osExit(1)
		return
	}
}

func runPut(cmd *cobra.Command, args []string) {
	if len(args) == 0 || len(args)%2 == 1 {
		cmd.Usage()
		return
	}

	var b client.Batch
	for i := 0; i < len(args); i += 2 {
		b.Put(
			unquoteArg(args[i], true /* disallow system keys */),
			unquoteArg(args[i+1], false),
		)
	}

	kvDB := makeDBClient()
	if kvDB == nil {
		return
	}

	if err := kvDB.Run(&b); err != nil {
		fmt.Fprintf(osStderr, "put failed: %s\n", err)
		osExit(1)
		return
	}
}

func (p *planner) makeBackfillBatch(tableName *parser.QualifiedName, tableDesc *TableDescriptor, indexDescs ...IndexDescriptor) (client.Batch, error) {
	b := client.Batch{}
	// Get all the rows affected.
	// TODO(vivek): Avoid going through Select.
	// TODO(tamird): Support partial indexes?
	row, err := p.Select(&parser.Select{
		Exprs: parser.SelectExprs{parser.StarSelectExpr()},
		From:  parser.TableExprs{&parser.AliasedTableExpr{Expr: tableName}},
	})
	if err != nil {
		return b, err
	}

	// Construct a map from column ID to the index the value appears at within a
	// row.
	colIDtoRowIndex := map[ColumnID]int{}
	for i, name := range row.Columns() {
		c, err := tableDesc.FindColumnByName(name)
		if err != nil {
			return b, err
		}
		colIDtoRowIndex[c.ID] = i
	}

	// TODO(tamird): This will fall down in production use. We need to do
	// something better (see #2036). In particular, this implementation
	// has the following problems:
	// - Very large tables will generate an enormous batch here. This
	// isn't really a problem in itself except that it will exacerbate
	// the other issue:
	// - Any non-quiescent table that this runs against will end up with
	// an inconsistent index. This is because as inserts/updates continue
	// to roll in behind this operation's read front, the written index
	// will become incomplete/stale before it's written.

	for row.Next() {
		rowVals := row.Values()

		for _, indexDesc := range indexDescs {
			secondaryIndexEntries, err := encodeSecondaryIndexes(
				tableDesc.ID, []IndexDescriptor{indexDesc}, colIDtoRowIndex, rowVals)
			if err != nil {
				return b, err
			}

			for _, secondaryIndexEntry := range secondaryIndexEntries {
				if log.V(2) {
					log.Infof("CPut %s -> %v", prettyKey(secondaryIndexEntry.key, 0),
						secondaryIndexEntry.value)
				}
				b.CPut(secondaryIndexEntry.key, secondaryIndexEntry.value, nil)
			}
		}
	}

	return b, row.Err()
}

// Delete deletes rows from a table.
func (p *planner) Delete(n *parser.Delete) (planNode, error) {
	tableDesc, err := p.getAliasedTableDesc(n.Table)
	if err != nil {
		return nil, err
	}

	// TODO(tamird,pmattis): avoid going through Select to avoid encoding
	// and decoding keys. Also, avoiding Select may provide more
	// convenient access to index keys which we are not currently
	// deleting.
	node, err := p.Select(&parser.Select{
		Exprs: parser.SelectExprs{
			&parser.StarExpr{TableName: parser.QualifiedName{tableDesc.Name}},
		},
		From:  parser.TableExprs{n.Table},
		Where: n.Where,
	})
	if err != nil {
		return nil, err
	}

	colMap := map[uint32]int{}
	for i, name := range node.Columns() {
		c, err := tableDesc.FindColumnByName(name)
		if err != nil {
			return nil, err
		}
		colMap[c.ID] = i
	}

	index := tableDesc.Indexes[0]
	indexKey := encodeIndexKeyPrefix(tableDesc.ID, index.ID)

	b := client.Batch{}

	for node.Next() {
		if err := node.Err(); err != nil {
			return nil, err
		}

		// TODO(tamird/pmattis): delete the secondary indexes too
		primaryKey, err := encodeIndexKey(index, colMap, node.Values(), indexKey)
		if err != nil {
			return nil, err
		}
		rowStartKey := proto.Key(primaryKey)
		b.DelRange(rowStartKey, rowStartKey.PrefixEnd())
	}

	if err := p.db.Run(&b); err != nil {
		return nil, err
	}

	// TODO(tamird/pmattis): return the number of affected rows
	return &valuesNode{}, nil
}

// createDescriptor implements the DescriptorAccessor interface.
func (p *planner) createDescriptor(plainKey sqlbase.DescriptorKey, descriptor sqlbase.DescriptorProto, ifNotExists bool) (bool, error) {
	idKey := plainKey.Key()
	// Check whether idKey exists.
	gr, err := p.txn.Get(idKey)
	if err != nil {
		return false, err
	}

	if gr.Exists() {
		if ifNotExists {
			// Noop.
			return false, nil
		}
		// Key exists, but we don't want it to: error out.
		return false, fmt.Errorf("%s %q already exists", descriptor.TypeName(), plainKey.Name())
	}

	// Increment unique descriptor counter.
	if ir, err := p.txn.Inc(keys.DescIDGenerator, 1); err == nil {
		descriptor.SetID(sqlbase.ID(ir.ValueInt() - 1))
	} else {
		return false, err
	}

	// TODO(pmattis): The error currently returned below is likely going to be
	// difficult to interpret.
	//
	// TODO(pmattis): Need to handle if-not-exists here as well.
	//
	// TODO(pmattis): This is writing the namespace and descriptor table entries,
	// but not going through the normal INSERT logic and not performing a precise
	// mimicry. In particular, we're only writing a single key per table, while
	// perfect mimicry would involve writing a sentinel key for each row as well.
	descKey := sqlbase.MakeDescMetadataKey(descriptor.GetID())

	b := client.Batch{}
	descID := descriptor.GetID()
	descDesc := sqlbase.WrapDescriptor(descriptor)
	if log.V(2) {
		log.Infof("CPut %s -> %d", idKey, descID)
		log.Infof("CPut %s -> %s", descKey, descDesc)
	}
	b.CPut(idKey, descID, nil)
	b.CPut(descKey, descDesc, nil)

	p.setTestingVerifyMetadata(func(systemConfig config.SystemConfig) error {
		if err := expectDescriptorID(systemConfig, idKey, descID); err != nil {
			return err
		}
		return expectDescriptor(systemConfig, descKey, descDesc)
	})

	return true, p.txn.Run(&b)
}