Golang Txn.NewBatch方法代码示例

func restoreTableDesc(
	ctx context.Context,
	txn *client.Txn,
	database sqlbase.DatabaseDescriptor,
	table sqlbase.TableDescriptor,
) error {
	// Run getDescriptorID again to make sure the database hasn't been dropped
	// while we were importing.
	var err error
	if table.ParentID, err = getDescriptorID(txn, tableKey{name: database.Name}); err != nil {
		return err
	}
	tableIDKey := tableKey{parentID: table.ParentID, name: table.Name}.Key()
	tableDescKey := sqlbase.MakeDescMetadataKey(table.ID)

	// Check for an existing table.
	var existingDesc sqlbase.Descriptor
	existingIDKV, err := txn.Get(tableIDKey)
	if err != nil {
		return err
	}
	if existingIDKV.Value != nil {
		existingID, err := existingIDKV.Value.GetInt()
		if err != nil {
			return err
		}
		existingDescKV, err := txn.Get(sqlbase.MakeDescMetadataKey(sqlbase.ID(existingID)))
		if err != nil {
			return err
		}
		if err := existingDescKV.Value.GetProto(&existingDesc); err != nil {
			return err
		}
	}

	// Write the new descriptors. First the ID -> TableDescriptor for the new
	// table, then flip (or initialize) the name -> ID entry so any new queries
	// will use the new one. If there was an existing table, it can now be
	// cleaned up.
	b := txn.NewBatch()
	b.CPut(tableDescKey, sqlbase.WrapDescriptor(&table), nil)
	if existingTable := existingDesc.GetTable(); existingTable == nil {
		b.CPut(tableIDKey, table.ID, nil)
	} else {
		existingIDKV.Value.ClearChecksum()
		b.CPut(tableIDKey, table.ID, existingIDKV.Value)
		// TODO(dan): This doesn't work for interleaved tables. Fix it when we
		// fix the empty range interleaved table TODO below.
		existingDataPrefix := roachpb.Key(keys.MakeTablePrefix(uint32(existingTable.ID)))
		b.DelRange(existingDataPrefix, existingDataPrefix.PrefixEnd(), false)
		zoneKey, _, descKey := GetKeysForTableDescriptor(existingTable)
		// Delete the desc and zone entries. Leave the name because the new
		// table is using it.
		b.Del(descKey)
		b.Del(zoneKey)
	}
	return txn.Run(b)
}

// Ingest loads some data in an sstable into an empty range. Only the keys
// between startKey and endKey are loaded. If newTableID is non-zero, every
// row's key is rewritten to be for that table.
func Ingest(
	ctx context.Context,
	txn *client.Txn,
	path string,
	checksum uint32,
	startKey, endKey roachpb.Key,
	newTableID sqlbase.ID,
) error {
	// TODO(mjibson): An appropriate value for this should be determined. The
	// current value was guessed at but appears to work well.
	const batchSize = 10000

	// TODO(dan): Check if the range being ingested into is empty. If newTableID
	// is non-zero, it'll have to be derived from startKey and endKey.

	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	crc := crc32.New(crc32.MakeTable(crc32.Castagnoli))
	if _, err := io.Copy(crc, f); err != nil {
		return nil
	}
	if c := crc.Sum32(); c != checksum {
		return errors.Errorf("%s: checksum mismatch got %d expected %d", path, c, checksum)
	}

	sst, err := engine.MakeRocksDBSstFileReader()
	if err != nil {
		return err
	}
	defer sst.Close()
	if err := sst.AddFile(path); err != nil {
		return err
	}

	b := txn.NewBatch()
	var v roachpb.Value
	count := 0
	ingestFunc := func(kv engine.MVCCKeyValue) (bool, error) {
		v = roachpb.Value{RawBytes: kv.Value}
		v.ClearChecksum()
		if log.V(3) {
			log.Infof(ctx, "Put %s %s\n", kv.Key.Key, v.PrettyPrint())
		}
		b.Put(kv.Key.Key, &v)
		count++
		if count > batchSize {
			if err := txn.Run(b); err != nil {
				return true, err
			}
			b = txn.NewBatch()
			count = 0
		}
		return false, nil
	}
	if newTableID != 0 {
		// MakeRekeyMVCCKeyValFunc modifies the keys, but this is safe because
		// the one we get back from rocksDBIterator.Key is a copy (not a
		// reference to the mmaped file.)
		ingestFunc = MakeRekeyMVCCKeyValFunc(newTableID, ingestFunc)
	}
	startKeyMVCC, endKeyMVCC := engine.MVCCKey{Key: startKey}, engine.MVCCKey{Key: endKey}
	if err := sst.Iterate(startKeyMVCC, endKeyMVCC, ingestFunc); err != nil {
		return err
	}
	return txn.Run(b)
}

func (td *tableDeleter) init(txn *client.Txn) error {
	td.txn = txn
	td.b = txn.NewBatch()
	return nil
}