Golang Transaction.Seek方法代碼示例

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

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

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

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

// Next implements plan.Plan Next interface.
func (r *TableDefaultPlan) Next(ctx context.Context) (row *plan.Row, err error) {
	if r.rangeScan {
		return r.rangeNext(ctx)
	}
	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(r.T.FirstKey())
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	if !r.iter.Valid() || !r.iter.Key().HasPrefix(r.T.RecordPrefix()) {
		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 := tables.DecodeRecordKeyHandle(rowKey)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// 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: string(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
}

func checkSeek(c *C, txn kv.Transaction) {
	for i := startIndex; i < testCount; i++ {
		val := encodeInt(i)
		iter, err := txn.Seek(val, nil)
		c.Assert(err, IsNil)
		c.Assert(iter.Key(), Equals, string(val))
		c.Assert(decodeInt([]byte(valToStr(c, iter))), Equals, i)
		iter.Close()
	}

	// Test iterator Next()
	for i := startIndex; i < testCount-1; i++ {
		val := encodeInt(i)
		iter, err := txn.Seek(val, nil)
		c.Assert(err, IsNil)
		c.Assert(iter.Key(), Equals, string(val))
		c.Assert(valToStr(c, iter), Equals, string(val))

		next, err := iter.Next(nil)
		c.Assert(err, IsNil)
		c.Assert(next.Valid(), IsTrue)

		val = encodeInt(i + 1)
		c.Assert(next.Key(), Equals, string(val))
		c.Assert(valToStr(c, next), Equals, string(val))
		iter.Close()
	}

	// Non exist seek test
	iter, err := txn.Seek(encodeInt(testCount), nil)
	c.Assert(err, IsNil)
	c.Assert(iter.Valid(), IsFalse)
	iter.Close()
}

func checkSeek(c *C, txn kv.Transaction) {
	for i := startIndex; i < testCount; i++ {
		val := encodeInt(i * indexStep)
		iter, err := txn.Seek(val)
		c.Assert(err, IsNil)
		c.Assert([]byte(iter.Key()), BytesEquals, val)
		c.Assert(decodeInt([]byte(valToStr(c, iter))), Equals, i*indexStep)
		iter.Close()
	}

	// Test iterator Next()
	for i := startIndex; i < testCount-1; i++ {
		val := encodeInt(i * indexStep)
		iter, err := txn.Seek(val)
		c.Assert(err, IsNil)
		c.Assert([]byte(iter.Key()), BytesEquals, val)
		c.Assert(valToStr(c, iter), Equals, string(val))

		err = iter.Next()
		c.Assert(err, IsNil)
		c.Assert(iter.Valid(), IsTrue)

		val = encodeInt((i + 1) * indexStep)
		c.Assert([]byte(iter.Key()), BytesEquals, val)
		c.Assert(valToStr(c, iter), Equals, string(val))
		iter.Close()
	}

	// Non exist and beyond maximum seek test
	iter, err := txn.Seek(encodeInt(testCount * indexStep))
	c.Assert(err, IsNil)
	c.Assert(iter.Valid(), IsFalse)

	// Non exist but between existing keys seek test,
	// it returns the smallest key that larger than the one we are seeking
	inBetween := encodeInt((testCount-1)*indexStep - 1)
	last := encodeInt((testCount - 1) * indexStep)
	iter, err = txn.Seek(inBetween)
	c.Assert(err, IsNil)
	c.Assert(iter.Valid(), IsTrue)
	c.Assert([]byte(iter.Key()), Not(BytesEquals), inBetween)
	c.Assert([]byte(iter.Key()), BytesEquals, last)
	iter.Close()
}

func getIndexRowFromRange(idxInfo *tipb.IndexInfo, txn kv.Transaction, ran kv.KeyRange, desc bool, limit int64) ([]*tipb.Row, error) {
	var rows []*tipb.Row
	var seekKey kv.Key
	if desc {
		seekKey = ran.EndKey
	} else {
		seekKey = ran.StartKey
	}
	for {
		if limit == 0 {
			break
		}
		var it kv.Iterator
		var err error
		if desc {
			it, err = txn.SeekReverse(seekKey)
			if err != nil {
				return nil, errors.Trace(err)
			}
			seekKey = it.Key()
		} else {
			it, err = txn.Seek(seekKey)
			if err != nil {
				return nil, errors.Trace(err)
			}
			seekKey = it.Key().PrefixNext()
		}
		if !it.Valid() {
			break
		}
		if desc {
			if it.Key().Cmp(ran.StartKey) < 0 {
				break
			}
		} else {
			if it.Key().Cmp(ran.EndKey) >= 0 {
				break
			}
		}

		datums, err := tablecodec.DecodeIndexKey(it.Key())
		if err != nil {
			return nil, errors.Trace(err)
		}
		var handle types.Datum
		if len(datums) > len(idxInfo.Columns) {
			handle = datums[len(idxInfo.Columns)]
			datums = datums[:len(idxInfo.Columns)]
		} else {
			var intHandle int64
			intHandle, err = decodeHandle(it.Value())
			if err != nil {
				return nil, errors.Trace(err)
			}
			handle.SetInt64(intHandle)
		}
		data, err := codec.EncodeValue(nil, datums...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		handleData, err := codec.EncodeValue(nil, handle)
		if err != nil {
			return nil, errors.Trace(err)
		}
		row := &tipb.Row{Handle: handleData, Data: data}
		rows = append(rows, row)
		limit--
	}

	return rows, nil
}