Golang kv.NextUntil函数代码示例

// Do implements the plan.Plan interface, iterates rows but does nothing.
func (r *TableNilPlan) Do(ctx context.Context, f plan.RowIterFunc) error {
	h := r.T.FirstKey()
	prefix := r.T.KeyPrefix()
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return err
	}
	it, err := txn.Seek([]byte(h), nil)
	if err != nil {
		return err
	}
	defer it.Close()
	for it.Valid() && strings.HasPrefix(it.Key(), prefix) {
		var err error
		id, err := util.DecodeHandleFromRowKey(it.Key())
		if err != nil {
			return err
		}

		// do nothing
		if m, err := f(id, nil); !m || err != nil {
			return err
		}

		rk := r.T.RecordKey(id, nil)
		if it, err = kv.NextUntil(it, util.RowKeyPrefixFilter(rk)); err != nil {
			return err
		}
	}
	return nil
}

// 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 updateOldRows(ctx context.Context, t *tables.Table, col *column.Col) error {
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return errors.Trace(err)
	}
	it, err := txn.Seek([]byte(t.FirstKey()))
	if err != nil {
		return errors.Trace(err)
	}
	defer it.Close()

	prefix := t.KeyPrefix()
	for it.Valid() && strings.HasPrefix(it.Key(), prefix) {
		handle, err0 := util.DecodeHandleFromRowKey(it.Key())
		if err0 != nil {
			return errors.Trace(err0)
		}
		k := t.RecordKey(handle, col)

		// TODO: check and get timestamp/datetime default value.
		// refer to getDefaultValue in stmt/stmts/stmt_helper.go.
		if err0 = t.SetColValue(txn, k, col.DefaultValue); err0 != nil {
			return errors.Trace(err0)
		}

		rk := t.RecordKey(handle, nil)
		if it, err0 = kv.NextUntil(it, util.RowKeyPrefixFilter(rk)); err0 != nil {
			return errors.Trace(err0)
		}
	}

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

// IterRecords implements table.Table IterRecords interface.
func (t *Table) IterRecords(ctx context.Context, startKey kv.Key, cols []*table.Column,
	fn table.RecordIterFunc) error {
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return errors.Trace(err)
	}
	it, err := txn.Seek(startKey)
	if err != nil {
		return errors.Trace(err)
	}
	defer it.Close()

	if !it.Valid() {
		return nil
	}

	log.Debugf("startKey:%q, key:%q, value:%q", startKey, it.Key(), it.Value())

	colMap := make(map[int64]*types.FieldType)
	for _, col := range cols {
		colMap[col.ID] = &col.FieldType
	}
	prefix := t.RecordPrefix()
	for it.Valid() && it.Key().HasPrefix(prefix) {
		// first kv pair is row lock information.
		// TODO: check valid lock
		// get row handle
		handle, err := tablecodec.DecodeRowKey(it.Key())
		if err != nil {
			return errors.Trace(err)
		}
		rowMap, err := tablecodec.DecodeRow(it.Value(), colMap)
		if err != nil {
			return errors.Trace(err)
		}
		data := make([]types.Datum, 0, len(cols))
		for _, col := range cols {
			if col.IsPKHandleColumn(t.Meta()) {
				data = append(data, types.NewIntDatum(handle))
			} else {
				data = append(data, rowMap[col.ID])
			}
		}
		more, err := fn(handle, data, cols)
		if !more || err != nil {
			return errors.Trace(err)
		}

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

	return nil
}

// Do scans over rows' kv pair in the table, and constructs them into row data.
func (r *TableDefaultPlan) Do(ctx context.Context, f plan.RowIterFunc) error {
	t := r.T
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return err
	}
	head := t.FirstKey()
	prefix := t.KeyPrefix()

	it, err := txn.Seek([]byte(head), nil)
	if err != nil {
		return err
	}
	defer it.Close()
	for it.Valid() && strings.HasPrefix(it.Key(), prefix) {
		// 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
		// ...
		var err error
		rowKey := it.Key()
		h, err := util.DecodeHandleFromRowKey(rowKey)
		if err != nil {
			return err
		}

		// TODO: we could just fetch mentioned columns' values
		rec, err := t.Row(ctx, h)
		if err != nil {
			return err
		}
		// Put rowKey to the tail of record row
		rks := &RowKeyList{}
		rke := &RowKeyEntry{
			Tbl: t,
			Key: rowKey,
		}
		rks.appendKeys(rke)
		rec = append(rec, rks)
		m, err := f(int64(0), rec)
		if !m || err != nil {
			return err
		}

		rk := t.RecordKey(h, nil)
		it, err = kv.NextUntil(it, util.RowKeyPrefixFilter(rk))
		if err != nil {
			return err
		}
	}
	return nil
}

func (d *ddl) buildIndex(ctx context.Context, t table.Table, idxInfo *model.IndexInfo, unique bool) error {
	firstKey := t.FirstKey()
	prefix := t.KeyPrefix()

	txn, err := ctx.GetTxn(false)
	if err != nil {
		return errors.Trace(err)
	}
	it, err := txn.Seek([]byte(firstKey))
	if err != nil {
		return errors.Trace(err)
	}
	defer it.Close()
	for it.Valid() && strings.HasPrefix(it.Key(), prefix) {
		var err error
		handle, err := util.DecodeHandleFromRowKey(it.Key())
		log.Info("building index...", handle)
		if err != nil {
			return errors.Trace(err)
		}
		// TODO: v is timestamp ?
		// fetch datas
		cols := t.Cols()
		var vals []interface{}
		for _, v := range idxInfo.Columns {
			var (
				data []byte
				val  interface{}
			)
			col := cols[v.Offset]
			k := t.RecordKey(handle, col)
			data, err = txn.Get([]byte(k))
			if err != nil {
				return errors.Trace(err)
			}
			val, err = t.DecodeValue(data, col)
			if err != nil {
				return errors.Trace(err)
			}
			vals = append(vals, val)
		}
		// build index
		kvX := kv.NewKVIndex(t.IndexPrefix(), idxInfo.Name.L, unique)
		err = kvX.Create(txn, vals, handle)
		if err != nil {
			return errors.Trace(err)
		}

		rk := []byte(t.RecordKey(handle, nil))
		it, err = kv.NextUntil(it, util.RowKeyPrefixFilter(rk))
		if err != nil {
			return errors.Trace(err)
		}
	}
	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 (d *ddl) getSnapshotRows(t table.Table, version uint64, seekHandle int64) ([]int64, error) {
	ver := kv.Version{Ver: version}

	snap, err := d.store.GetSnapshot(ver)
	if err != nil {
		return nil, errors.Trace(err)
	}

	defer snap.Release()

	firstKey := t.RecordKey(seekHandle, nil)

	it, err := snap.Seek(firstKey)
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer it.Close()

	handles := make([]int64, 0, maxBatchSize)

	for it.Valid() {
		if !it.Key().HasPrefix(t.RecordPrefix()) {
			break
		}

		var handle int64
		handle, err = tables.DecodeRecordKeyHandle(it.Key())
		if err != nil {
			return nil, errors.Trace(err)
		}

		rk := t.RecordKey(handle, nil)

		handles = append(handles, handle)
		if len(handles) == maxBatchSize {
			seekHandle = handle + 1
			break
		}

		err = kv.NextUntil(it, util.RowKeyPrefixFilter(rk))
		if terror.ErrorEqual(err, kv.ErrNotExist) {
			break
		} else if err != nil {
			return nil, errors.Trace(err)
		}
	}

	return handles, nil
}

func (d *ddl) getSnapshotRows(t table.Table, version uint64, seekHandle int64) ([]int64, error) {
	ver := kv.Version{Ver: version}

	snap, err := d.store.GetSnapshot(ver)
	if err != nil {
		return nil, errors.Trace(err)
	}

	defer snap.MvccRelease()

	firstKey := t.RecordKey(seekHandle, nil)
	prefix := []byte(t.KeyPrefix())

	it := snap.NewMvccIterator([]byte(firstKey), ver)
	defer it.Close()

	handles := make([]int64, 0, maxBatchSize)

	for it.Valid() {
		key := []byte(it.Key())
		if !bytes.HasPrefix(key, prefix) {
			break
		}

		var handle int64
		handle, err = util.DecodeHandleFromRowKey(string(key))
		if err != nil {
			return nil, errors.Trace(err)
		}

		rk := t.RecordKey(handle, nil)

		handles = append(handles, handle)
		if len(handles) == maxBatchSize {
			seekHandle = handle + 1
			break
		}

		err = kv.NextUntil(it, util.RowKeyPrefixFilter(rk))
		if terror.ErrorEqual(err, kv.ErrNotExist) {
			break
		} else if err != nil {
			return nil, errors.Trace(err)
		}
	}

	return handles, nil
}

// IterRecords implements table.Table IterRecords interface.
func (t *Table) IterRecords(ctx context.Context, startKey string, cols []*column.Col, fn table.RecordIterFunc) error {
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return errors.Trace(err)
	}

	it, err := txn.Seek([]byte(startKey))
	if err != nil {
		return errors.Trace(err)
	}
	defer it.Close()

	if !it.Valid() {
		return nil
	}

	log.Debugf("startKey %q, key:%q,value:%q", startKey, it.Key(), it.Value())

	prefix := t.KeyPrefix()
	for it.Valid() && strings.HasPrefix(it.Key(), prefix) {
		// first kv pair is row lock information.
		// TODO: check valid lock
		// get row handle
		var err error
		handle, err := util.DecodeHandleFromRowKey(it.Key())
		if err != nil {
			return errors.Trace(err)
		}

		data, err := t.RowWithCols(ctx, handle, cols)
		if err != nil {
			return errors.Trace(err)
		}
		more, err := fn(handle, data, cols)
		if !more || err != nil {
			return errors.Trace(err)
		}

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

func iterRecords(retriever kv.Retriever, t table.Table, startKey kv.Key, cols []*table.Column,
	fn table.RecordIterFunc) error {
	it, err := retriever.Seek(startKey)
	if err != nil {
		return errors.Trace(err)
	}
	defer it.Close()

	if !it.Valid() {
		return nil
	}

	log.Debugf("startKey:%q, key:%q, value:%q", startKey, it.Key(), it.Value())

	prefix := t.RecordPrefix()
	for it.Valid() && it.Key().HasPrefix(prefix) {
		// first kv pair is row lock information.
		// TODO: check valid lock
		// get row handle
		handle, err := tables.DecodeRecordKeyHandle(it.Key())
		if err != nil {
			return errors.Trace(err)
		}

		data, err := rowWithCols(retriever, t, handle, cols)
		if err != nil {
			return errors.Trace(err)
		}
		more, err := fn(handle, data, cols)
		if !more || err != nil {
			return errors.Trace(err)
		}

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

	return nil
}