Golang kv.Transaction类代码示例

// backfillIndexInTxn deals with a part of backfilling index data in a Transaction.
// This part of the index data rows is defaultSmallBatchCnt.
func (d *ddl) backfillIndexInTxn(t table.Table, kvIdx table.Index, handles []int64, txn kv.Transaction) (int64, error) {
	idxRecords, err := d.fetchRowColVals(txn, t, handles, kvIdx.Meta())
	if err != nil {
		return 0, errors.Trace(err)
	}

	for _, idxRecord := range idxRecords {
		log.Debug("[ddl] backfill index...", idxRecord.handle)
		err = txn.LockKeys(idxRecord.key)
		if err != nil {
			return 0, errors.Trace(err)
		}

		// Create the index.
		handle, err := kvIdx.Create(txn, idxRecord.vals, idxRecord.handle)
		if err != nil {
			if terror.ErrorEqual(err, kv.ErrKeyExists) && idxRecord.handle == handle {
				// Index already exists, skip it.
				continue
			}
			return 0, errors.Trace(err)
		}
	}
	return idxRecords[len(idxRecords)-1].handle, nil
}

// Next implements plan.Plan Next interface.
func (r *TableDefaultPlan) Next(ctx context.Context) (row *plan.Row, err error) {
	if r.iter == nil {
		var txn kv.Transaction
		txn, err = ctx.GetTxn(false)
		if err != nil {
			return nil, errors.Trace(err)
		}
		r.iter, err = txn.Seek([]byte(r.T.FirstKey()))
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	if !r.iter.Valid() || !strings.HasPrefix(r.iter.Key(), r.T.KeyPrefix()) {
		return
	}
	// TODO: check if lock valid
	// the record layout in storage (key -> value):
	// r1 -> lock-version
	// r1_col1 -> r1 col1 value
	// r1_col2 -> r1 col2 value
	// r2 -> lock-version
	// r2_col1 -> r2 col1 value
	// r2_col2 -> r2 col2 value
	// ...
	rowKey := r.iter.Key()
	handle, err := util.DecodeHandleFromRowKey(rowKey)
	if err != nil {
		return nil, errors.Trace(err)
	}

	txn, err := ctx.GetTxn(false)
	if err != nil {
		return nil, errors.Trace(err)
	}
	// It is very likely that we will fetch rows after current row later, enable the RangePrefetchOnCacheMiss
	// option may help reducing RPC calls.
	// TODO: choose a wiser option value.
	txn.SetOption(kv.RangePrefetchOnCacheMiss, nil)
	defer txn.DelOption(kv.RangePrefetchOnCacheMiss)

	// TODO: we could just fetch mentioned columns' values
	row = &plan.Row{}
	row.Data, err = r.T.Row(ctx, handle)
	if err != nil {
		return nil, errors.Trace(err)
	}
	// Put rowKey to the tail of record row
	rke := &plan.RowKeyEntry{
		Tbl: r.T,
		Key: rowKey,
	}
	row.RowKeys = append(row.RowKeys, rke)

	rk := r.T.RecordKey(handle, nil)
	err = kv.NextUntil(r.iter, util.RowKeyPrefixFilter(rk))
	if err != nil {
		return nil, errors.Trace(err)
	}
	return
}

// GetTableRecordsCount returns the total number of table records from startHandle.
// If startHandle = 0, returns the total number of table records.
func GetTableRecordsCount(txn kv.Transaction, t table.Table, startHandle int64) (int64, error) {
	startKey := t.RecordKey(startHandle, nil)
	it, err := txn.Seek(startKey)
	if err != nil {
		return 0, errors.Trace(err)
	}

	var cnt int64
	prefix := t.RecordPrefix()
	for it.Valid() && it.Key().HasPrefix(prefix) {
		handle, err := tables.DecodeRecordKeyHandle(it.Key())
		if err != nil {
			return 0, errors.Trace(err)
		}

		it.Close()
		rk := t.RecordKey(handle+1, nil)
		it, err = txn.Seek(rk)
		if err != nil {
			return 0, errors.Trace(err)
		}

		cnt++
	}

	it.Close()

	return cnt, nil
}

func lockRow(txn kv.Transaction, t table.Table, h int64) error {
	// Get row lock key
	lockKey := t.RecordKey(h, nil)
	// set row lock key to current txn
	err := txn.Set(lockKey, []byte(txn.String()))
	return errors.Trace(err)
}

// ScanMetaWithPrefix scans metadata with the prefix.
func ScanMetaWithPrefix(txn kv.Transaction, prefix string, filter func([]byte, []byte) bool) error {
	iter, err := txn.Seek([]byte(prefix))
	if err != nil {
		return errors.Trace(err)
	}
	defer iter.Close()

	for {
		if err != nil {
			return errors.Trace(err)
		}

		if iter.Valid() && strings.HasPrefix(iter.Key(), prefix) {
			if !filter([]byte(iter.Key()), iter.Value()) {
				break
			}
			err = iter.Next()
			if err != nil {
				return errors.Trace(err)
			}
		} else {
			break
		}
	}

	return nil
}

func mustNotGet(c *C, txn kv.Transaction) {
	for i := startIndex; i < testCount; i++ {
		s := encodeInt(i * indexStep)
		_, err := txn.Get(s)
		c.Assert(err, NotNil)
	}
}

func fetchRowColVals(txn kv.Transaction, t table.Table, handle int64, indexInfo *model.IndexInfo) (
	kv.Key, []types.Datum, error) {
	// fetch datas
	cols := t.Cols()
	colMap := make(map[int64]*types.FieldType)
	for _, v := range indexInfo.Columns {
		col := cols[v.Offset]
		colMap[col.ID] = &col.FieldType
	}
	rowKey := tablecodec.EncodeRecordKey(t.RecordPrefix(), handle)
	rowVal, err := txn.Get(rowKey)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	row, err := tablecodec.DecodeRow(rowVal, colMap)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	vals := make([]types.Datum, 0, len(indexInfo.Columns))
	for _, v := range indexInfo.Columns {
		col := cols[v.Offset]
		vals = append(vals, row[col.ID])
	}
	return rowKey, vals, nil
}

func insertData(c *C, txn kv.Transaction) {
	for i := startIndex; i < testCount; i++ {
		val := encodeInt(i * indexStep)
		err := txn.Set(val, val)
		c.Assert(err, IsNil)
	}
}

// Next implements plan.Plan Next interface.
func (r *TableNilPlan) Next(ctx context.Context) (row *plan.Row, err error) {
	if r.iter == nil {
		var txn kv.Transaction
		txn, err = ctx.GetTxn(false)
		if err != nil {
			return nil, errors.Trace(err)
		}
		r.iter, err = txn.Seek([]byte(r.T.FirstKey()), nil)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	if !r.iter.Valid() || !strings.HasPrefix(r.iter.Key(), r.T.KeyPrefix()) {
		return
	}
	id, err := util.DecodeHandleFromRowKey(r.iter.Key())
	if err != nil {
		return nil, errors.Trace(err)
	}
	rk := r.T.RecordKey(id, nil)
	// Even though the data is nil, we should return not nil row,
	// or the iteration will stop.
	row = &plan.Row{}
	r.iter, err = kv.NextUntil(r.iter, util.RowKeyPrefixFilter(rk))
	return
}

func setRow(txn kv.Transaction, handle int64, tbl *simpleTableInfo, gen genValueFunc) error {
	rowKey := tablecodec.EncodeRowKey(tbl.tID, codec.EncodeInt(nil, handle))
	columnValues := gen(handle, tbl)
	value, err := tablecodec.EncodeRow(columnValues, tbl.cIDs)
	if err != nil {
		return errors.Trace(err)
	}
	err = txn.Set(rowKey, value)
	if err != nil {
		return errors.Trace(err)
	}
	for i, idxCol := range tbl.indices {
		idxVal := columnValues[idxCol]
		encoded, err := codec.EncodeKey(nil, idxVal, types.NewDatum(handle))
		if err != nil {
			return errors.Trace(err)
		}
		idxKey := tablecodec.EncodeIndexSeekKey(tbl.tID, tbl.iIDs[i], encoded)
		err = txn.Set(idxKey, []byte{0})
		if err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func mustDel(c *C, txn kv.Transaction) {
	for i := startIndex; i < testCount; i++ {
		val := encodeInt(i * indexStep)
		err := txn.Delete(val)
		c.Assert(err, IsNil)
	}
}

// SetDDLBinlog sets DDL binlog in the kv.Transaction.
func SetDDLBinlog(txn kv.Transaction, jobID int64, ddlQuery string) {
	bin := &binlog.Binlog{
		Tp:       binlog.BinlogType_Prewrite,
		DdlJobId: jobID,
		DdlQuery: []byte(ddlQuery),
	}
	txn.SetOption(kv.BinlogData, bin)
}

func setColValue(c *C, txn kv.Transaction, key kv.Key, v types.Datum) {
	row := []types.Datum{v, {}}
	colIDs := []int64{2, 3}
	value, err := tablecodec.EncodeRow(row, colIDs)
	c.Assert(err, IsNil)
	err = txn.Set(key, value)
	c.Assert(err, IsNil)
}

func mustGet(c *C, txn kv.Transaction) {
	for i := startIndex; i < testCount; i++ {
		s := encodeInt(i * indexStep)
		val, err := txn.Get(s)
		c.Assert(err, IsNil)
		c.Assert(string(val), Equals, string(s))
	}
}

// SetColValue implements table.Table SetColValue interface.
func (t *Table) SetColValue(txn kv.Transaction, key []byte, data interface{}) error {
	v, err := t.EncodeValue(data)
	if err != nil {
		return errors.Trace(err)
	}
	if err := txn.Set(key, v); err != nil {
		return errors.Trace(err)
	}
	return nil
}