Golang Batch.Put方法代码示例

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
	}
}

// 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
}

// 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
}

// TestAuthentication tests authentication for the KV endpoint.
func TestAuthentication(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s := server.StartTestServer(t)
	defer s.Stop()

	var b1 client.Batch
	b1.Put("a", "b")

	// Create a node user client and call Run() on it which lets us build our own
	// request, specifying the user.
	db1 := createTestClientForUser(t, s.Stopper(), s.ServingAddr(), security.NodeUser)
	if err := db1.Run(&b1); err != nil {
		t.Fatal(err)
	}

	var b2 client.Batch
	b2.Put("c", "d")

	// Try again, but this time with certs for a non-node user (even the root
	// user has no KV permissions).
	db2 := createTestClientForUser(t, s.Stopper(), s.ServingAddr(), security.RootUser)
	if err := db2.Run(&b2); !testutils.IsError(err, "is not allowed") {
		t.Fatal(err)
	}
}

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
	}
}

// RenameDatabase renames the database.
// Privileges: security.RootUser user.
//   Notes: postgres requires superuser, db owner, or "CREATEDB".
//          mysql >= 5.1.23 does not allow database renames.
func (p *planner) RenameDatabase(n *parser.RenameDatabase) (planNode, error) {
	if n.Name == "" || n.NewName == "" {
		return nil, errEmptyDatabaseName
	}

	if p.session.User != security.RootUser {
		return nil, fmt.Errorf("only %s is allowed to rename databases", security.RootUser)
	}

	dbDesc, err := p.getDatabaseDesc(string(n.Name))
	if err != nil {
		return nil, err
	}
	if dbDesc == nil {
		return nil, databaseDoesNotExistError(string(n.Name))
	}

	if n.Name == n.NewName {
		// Noop.
		return &emptyNode{}, nil
	}

	// Now update the nameMetadataKey and the descriptor.
	descKey := sqlbase.MakeDescMetadataKey(dbDesc.GetID())
	dbDesc.SetName(string(n.NewName))

	if err := dbDesc.Validate(); err != nil {
		return nil, err
	}

	newKey := databaseKey{string(n.NewName)}.Key()
	oldKey := databaseKey{string(n.Name)}.Key()
	descID := dbDesc.GetID()
	descDesc := sqlbase.WrapDescriptor(dbDesc)

	b := client.Batch{}
	b.CPut(newKey, descID, nil)
	b.Put(descKey, descDesc)
	b.Del(oldKey)

	if err := p.txn.Run(&b); err != nil {
		if _, ok := err.(*roachpb.ConditionFailedError); ok {
			return nil, fmt.Errorf("the new database name %q already exists", string(n.NewName))
		}
		return nil, err
	}

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

	return &emptyNode{}, nil
}

// flush writes all dirty nodes and the tree to the transaction.
func (tc *treeContext) flush(b *client.Batch) error {
	if tc.dirty {
		b.Put(keys.RangeTreeRoot, tc.tree)
	}
	for _, cachedNode := range tc.nodes {
		if cachedNode.dirty {
			b.Put(keys.RangeTreeNodeKey(cachedNode.node.Key), cachedNode.node)
		}
	}
	return nil
}

// CreateIndex creates an index.
// Privileges: CREATE on table.
//   notes: postgres requires CREATE on the table.
//          mysql requires INDEX on the table.
func (p *planner) CreateIndex(n *parser.CreateIndex) (planNode, error) {
	tableDesc, err := p.getTableDesc(n.Table)
	if err != nil {
		return nil, err
	}

	if _, err := tableDesc.FindIndexByName(string(n.Name)); err == nil {
		if n.IfNotExists {
			// Noop.
			return &valuesNode{}, nil
		}
		return nil, fmt.Errorf("index %q already exists", string(n.Name))
	}

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

	indexDesc := IndexDescriptor{
		Name:             string(n.Name),
		Unique:           n.Unique,
		ColumnNames:      n.Columns,
		StoreColumnNames: n.Storing,
	}

	newTableDesc := proto.Clone(tableDesc).(*TableDescriptor)

	if err := newTableDesc.AddIndex(indexDesc, false); err != nil {
		return nil, err
	}

	if err := newTableDesc.AllocateIDs(); err != nil {
		return nil, err
	}

	b := client.Batch{}
	if err := p.backfillBatch(&b, n.Table, tableDesc, newTableDesc); err != nil {
		return nil, 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 nil, convertBatchError(newTableDesc, b, err)
	}

	return &valuesNode{}, nil
}

// flush writes all dirty nodes and the tree to the transaction.
func (tc *treeContext) flush(b *client.Batch) {
	if tc.dirty {
		b.Put(keys.RangeTreeRoot, tc.tree)
	}
	for key, cachedNode := range tc.nodes {
		if cachedNode.dirty {
			if cachedNode.node == nil {
				b.Del(keys.RangeTreeNodeKey(roachpb.RKey(key)))
			} else {
				b.Put(keys.RangeTreeNodeKey(roachpb.RKey(key)), cachedNode.node)
			}
		}
	}
}

// DropIndex drops an index.
// Privileges: CREATE on table.
//   Notes: postgres allows only the index owner to DROP an index.
//          mysql requires the INDEX privilege on the table.
func (p *planner) DropIndex(n *parser.DropIndex) (planNode, error) {
	b := client.Batch{}
	for _, indexQualifiedName := range n.Names {
		if err := indexQualifiedName.NormalizeTableName(p.session.Database); err != nil {
			return nil, err
		}

		tableDesc, err := p.getTableDesc(indexQualifiedName)
		if err != nil {
			return nil, err
		}

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

		newTableDesc := proto.Clone(tableDesc).(*TableDescriptor)

		idxName := indexQualifiedName.Index()
		i, err := newTableDesc.FindIndexByName(idxName)
		if err != nil {
			if n.IfExists {
				// Noop.
				return &valuesNode{}, nil
			}
			// Index does not exist, but we want it to: error out.
			return nil, err
		}
		newTableDesc.Indexes = append(newTableDesc.Indexes[:i], newTableDesc.Indexes[i+1:]...)

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

		if err := newTableDesc.Validate(); err != nil {
			return nil, err
		}

		descKey := MakeDescMetadataKey(newTableDesc.GetID())
		b.Put(descKey, wrapDescriptor(newTableDesc))
	}

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

	return &valuesNode{}, nil
}

// RenameDatabase renames the database.
// Privileges: "root" user.
//   Notes: postgres requires superuser, db owner, or "CREATEDB".
//          mysql >= 5.1.23 does not allow database renames.
func (p *planner) RenameDatabase(n *parser.RenameDatabase) (planNode, error) {
	if n.Name == "" || n.NewName == "" {
		return nil, errEmptyDatabaseName
	}

	if p.user != security.RootUser {
		return nil, fmt.Errorf("only %s is allowed to rename databases", security.RootUser)
	}

	dbDesc, err := p.getDatabaseDesc(string(n.Name))
	if err != nil {
		return nil, err
	}

	if n.Name == n.NewName {
		// Noop.
		return &valuesNode{}, nil
	}

	// Now update the nameMetadataKey and the descriptor.
	descKey := MakeDescMetadataKey(dbDesc.GetID())
	dbDesc.SetName(string(n.NewName))

	if err := dbDesc.Validate(); err != nil {
		return nil, err
	}

	b := client.Batch{}
	b.CPut(databaseKey{string(n.NewName)}.Key(), dbDesc.GetID(), nil)
	b.Put(descKey, dbDesc)
	b.Del(databaseKey{string(n.Name)}.Key())

	// Mark transaction as operating on the system DB.
	p.txn.SetSystemDBTrigger()
	if err := p.txn.Run(&b); err != nil {
		if _, ok := err.(*proto.ConditionFailedError); ok {
			return nil, fmt.Errorf("the new database name %q already exists", string(n.NewName))
		}
		return nil, err
	}

	return &valuesNode{}, nil
}

func runPut(cmd *cobra.Command, args []string) {
	if len(args) == 0 || len(args)%2 == 1 {
		mustUsage(cmd)
		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, stopper := makeDBClient()
	defer stopper.Stop()

	if err := kvDB.Run(&b); err != nil {
		panicf("put failed: %s", err)
	}
}

//.........这里部分代码省略.........
	// Secondary indexes needing updating.
	var indexes []IndexDescriptor
	for _, index := range tableDesc.Indexes {
		for _, id := range index.ColumnIDs {
			if _, ok := colIDSet[id]; ok {
				indexes = append(indexes, index)
				break
			}
		}
	}

	// Update all the rows.
	var b client.Batch
	for rows.Next() {
		rowVals := rows.Values()
		primaryIndexKey, _, err := encodeIndexKey(
			primaryIndex.ColumnIDs, colIDtoRowIndex, rowVals, primaryIndexKeyPrefix)
		if err != nil {
			return nil, err
		}
		// Compute the current secondary index key:value pairs for this row.
		secondaryIndexEntries, err := encodeSecondaryIndexes(
			tableDesc.ID, indexes, colIDtoRowIndex, rowVals)
		if err != nil {
			return nil, err
		}

		// Our updated value expressions occur immediately after the plain
		// columns in the output.
		newVals := rowVals[len(tableDesc.Columns):]
		// Update the row values.
		for i, col := range cols {
			val := newVals[i]
			if !col.Nullable && val == parser.DNull {
				return nil, fmt.Errorf("null value in column %q violates not-null constraint", col.Name)
			}
			rowVals[colIDtoRowIndex[col.ID]] = val
		}
		// Compute the new secondary index key:value pairs for this row.
		newSecondaryIndexEntries, err := encodeSecondaryIndexes(
			tableDesc.ID, indexes, colIDtoRowIndex, rowVals)
		if err != nil {
			return nil, err
		}
		// Update secondary indexes.
		for i, newSecondaryIndexEntry := range newSecondaryIndexEntries {
			secondaryIndexEntry := secondaryIndexEntries[i]
			if !bytes.Equal(newSecondaryIndexEntry.key, secondaryIndexEntry.key) {
				if log.V(2) {
					log.Infof("CPut %q -> %v", newSecondaryIndexEntry.key, newSecondaryIndexEntry.value)
				}
				b.CPut(newSecondaryIndexEntry.key, newSecondaryIndexEntry.value, nil)
				if log.V(2) {
					log.Infof("Del %q", secondaryIndexEntry.key)
				}
				b.Del(secondaryIndexEntry.key)
			}
		}

		// Add the new values.
		for i, val := range newVals {
			col := cols[i]

			primitive, err := convertDatum(col, val)
			if err != nil {
				return nil, err
			}

			key := MakeColumnKey(col.ID, primaryIndexKey)
			if primitive != nil {
				// We only output non-NULL values. Non-existent column keys are
				// considered NULL during scanning and the row sentinel ensures we know
				// the row exists.
				if log.V(2) {
					log.Infof("Put %q -> %v", key, val)
				}

				b.Put(key, primitive)
			} else {
				// The column might have already existed but is being set to NULL, so
				// delete it.
				if log.V(2) {
					log.Infof("Del %q", key)
				}

				b.Del(key)
			}
		}
	}

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

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

// updateRow adds to the batch the kv operations necessary to update a table row
// with the given values.
//
// The row corresponding to oldValues is updated with the ones in updateValues.
// Note that updateValues only contains the ones that are changing.
//
// The return value is only good until the next call to UpdateRow.
func (ru *rowUpdater) updateRow(
	b *client.Batch,
	oldValues []parser.Datum,
	updateValues []parser.Datum,
) ([]parser.Datum, error) {
	if len(oldValues) != len(ru.fetchCols) {
		return nil, util.Errorf("got %d values but expected %d", len(oldValues), len(ru.fetchCols))
	}
	if len(updateValues) != len(ru.updateCols) {
		return nil, util.Errorf("got %d values but expected %d", len(updateValues), len(ru.updateCols))
	}

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

	// Check that the new value types match the column types. This needs to
	// happen before index encoding because certain datum types (i.e. tuple)
	// cannot be used as index values.
	for i, val := range updateValues {
		if ru.marshalled[i], err = sqlbase.MarshalColumnValue(ru.updateCols[i], val); err != nil {
			return nil, err
		}
	}

	// Update the row values.
	copy(ru.newValues, oldValues)
	for i, updateCol := range ru.updateCols {
		ru.newValues[ru.fetchColIDtoRowIndex[updateCol.ID]] = updateValues[i]
	}

	newPrimaryIndexKey := primaryIndexKey
	rowPrimaryKeyChanged := false
	var newSecondaryIndexEntries []sqlbase.IndexEntry
	if ru.primaryKeyColChange {
		newPrimaryIndexKey, newSecondaryIndexEntries, err = ru.helper.encodeIndexes(ru.fetchColIDtoRowIndex, ru.newValues)
		if err != nil {
			return nil, err
		}
		rowPrimaryKeyChanged = !bytes.Equal(primaryIndexKey, newPrimaryIndexKey)
	} else {
		newSecondaryIndexEntries, err = sqlbase.EncodeSecondaryIndexes(
			ru.helper.tableDesc.ID, ru.helper.indexes, ru.fetchColIDtoRowIndex, ru.newValues)
		if err != nil {
			return nil, err
		}
	}

	if rowPrimaryKeyChanged {
		err := ru.rd.deleteRow(b, oldValues)
		if err != nil {
			return nil, err
		}
		err = ru.ri.insertRow(b, ru.newValues)
		return ru.newValues, err
	}

	// Update secondary indexes.
	for i, newSecondaryIndexEntry := range newSecondaryIndexEntries {
		secondaryIndexEntry := secondaryIndexEntries[i]
		secondaryKeyChanged := !bytes.Equal(newSecondaryIndexEntry.Key, secondaryIndexEntry.Key)
		if secondaryKeyChanged {
			if log.V(2) {
				log.Infof("Del %s", secondaryIndexEntry.Key)
			}
			b.Del(secondaryIndexEntry.Key)
			// Do not update Indexes in the DELETE_ONLY state.
			if _, ok := ru.deleteOnlyIndex[i]; !ok {
				if log.V(2) {
					log.Infof("CPut %s -> %v", newSecondaryIndexEntry.Key, newSecondaryIndexEntry.Value)
				}
				b.CPut(newSecondaryIndexEntry.Key, newSecondaryIndexEntry.Value, nil)
			}
		}
	}

	// Add the new values.
	for i, val := range updateValues {
		col := ru.updateCols[i]

		if ru.helper.columnInPK(col.ID) {
			// Skip primary key columns as their values are encoded in the row
			// sentinel key which is guaranteed to exist for as long as the row
			// exists.
			continue
		}

		ru.key = keys.MakeColumnKey(newPrimaryIndexKey, uint32(col.ID))
		if ru.marshalled[i].RawBytes != nil {
			// We only output non-NULL values. Non-existent column keys are
			// considered NULL during scanning and the row sentinel ensures we know
			// the row exists.
//.........这里部分代码省略.........

// RenameTable renames the table.
// Privileges: WRITE on database.
//   Notes: postgres requires the table owner.
//          mysql requires ALTER, DROP on the original table, and CREATE, INSERT
//          on the new table (and does not copy privileges over).
func (p *planner) RenameTable(n *parser.RenameTable) (planNode, error) {
	if n.NewName == "" {
		return nil, errEmptyTableName
	}

	if err := n.Name.NormalizeTableName(p.session.Database); err != nil {
		return nil, err
	}

	if n.Name.Table() == string(n.NewName) {
		// Noop.
		return &valuesNode{}, nil
	}

	dbDesc, err := p.getDatabaseDesc(n.Name.Database())
	if err != nil {
		return nil, err
	}

	tbKey := tableKey{dbDesc.ID, string(n.Name.Table())}.Key()

	// Check if table exists.
	gr, err := p.txn.Get(tbKey)
	if err != nil {
		return nil, err
	}
	if !gr.Exists() {
		if n.IfExists {
			// Noop.
			return &valuesNode{}, nil
		}
		// Key does not exist, but we want it to: error out.
		return nil, fmt.Errorf("table %q does not exist", n.Name.Table())
	}

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

	tableDesc, err := p.getTableDesc(n.Name)
	if err != nil {
		return nil, err
	}

	tableDesc.SetName(string(n.NewName))

	newTbKey := tableKey{dbDesc.ID, string(n.NewName)}.Key()
	descKey := MakeDescMetadataKey(tableDesc.GetID())

	b := client.Batch{}
	b.Put(descKey, tableDesc)
	b.CPut(newTbKey, descKey, nil)
	b.Del(tbKey)
	if err := p.txn.Run(&b); err != nil {
		if _, ok := err.(*proto.ConditionFailedError); ok {
			return nil, fmt.Errorf("table name %q already exists", n.NewName)
		}
		return nil, err
	}

	return &valuesNode{}, nil
}