Golang Context.GetTxn方法代码示例

// Do implements plan.Plan Do interface, acquiring locks.
func (r *SelectLockPlan) Do(ctx context.Context, f plan.RowIterFunc) error {
	return r.Src.Do(ctx, func(rid interface{}, in []interface{}) (bool, error) {
		var rowKeys *RowKeyList
		if in != nil && len(in) > 0 {
			t := in[len(in)-1]
			switch vt := t.(type) {
			case *RowKeyList:
				rowKeys = vt
				// Remove row key list from data tail
				in = in[:len(in)-1]
			}
		}
		if rowKeys != nil && r.Lock == coldef.SelectLockForUpdate {
			txn, err := ctx.GetTxn(false)
			if err != nil {
				return false, errors.Trace(err)
			}
			for _, k := range rowKeys.Keys {
				err = txn.LockKeys([]byte(k.Key))
				if err != nil {
					return false, errors.Trace(err)
				}
			}
		}
		return f(rid, in)
	})
}

// DelKeyWithPrefix deletes keys with prefix.
func DelKeyWithPrefix(ctx context.Context, prefix string) error {
	log.Debug("delKeyWithPrefix", prefix)
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return err
	}

	var keys []string
	iter, err := txn.Seek([]byte(prefix), hasPrefix([]byte(prefix)))
	if err != nil {
		return err
	}
	defer iter.Close()
	for {
		if err != nil {
			return err
		}

		if iter.Valid() && strings.HasPrefix(iter.Key(), prefix) {
			keys = append(keys, iter.Key())
			iter, err = iter.Next(hasPrefix([]byte(prefix)))
		} else {
			break
		}
	}

	for _, key := range keys {
		err := txn.Delete([]byte(key))
		if err != nil {
			return err
		}
	}

	return nil
}

// RowWithCols implements table.Table RowWithCols interface.
func (t *Table) RowWithCols(ctx context.Context, h int64, cols []*column.Col) ([]types.Datum, error) {
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return nil, errors.Trace(err)
	}
	v := make([]types.Datum, len(cols))
	for i, col := range cols {
		if col.State != model.StatePublic {
			return nil, errors.Errorf("Cannot use none public column - %v", cols)
		}
		if col.IsPKHandleColumn(t.meta) {
			v[i].SetInt64(h)
			continue
		}

		k := t.RecordKey(h, col)
		data, err := txn.Get(k)
		if err != nil {
			return nil, errors.Trace(err)
		}

		v[i], err = DecodeValue(data, &col.FieldType)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	return v, nil
}

// AddRecord implements table.Table AddRecord interface.
func (t *Table) AddRecord(ctx context.Context, r []interface{}) (recordID int64, err error) {
	id := variable.GetSessionVars(ctx).LastInsertID
	// Already have auto increment ID
	if id != 0 {
		recordID = int64(id)
	} else {
		recordID, err = t.alloc.Alloc(t.ID)
		if err != nil {
			return 0, err
		}
	}
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return 0, err
	}
	for _, v := range t.indices {
		if v == nil {
			continue
		}
		colVals, _ := v.FetchValues(r)
		if err = v.X.Create(txn, colVals, recordID); err != nil {
			if errors2.ErrorEqual(err, kv.ErrKeyExists) {
				// Get the duplicate row handle
				iter, _, terr := v.X.Seek(txn, colVals)
				if terr != nil {
					return 0, errors.Trace(terr)
				}
				_, h, terr := iter.Next()
				if terr != nil {
					return 0, errors.Trace(terr)
				}
				return h, errors.Trace(err)
			}
			return 0, errors.Trace(err)
		}
	}

	// split a record into multiple kv pair
	// first key -> LOCK
	k := t.RecordKey(recordID, nil)
	// A new row with current txn-id as lockKey
	err = txn.Set([]byte(k), []byte(txn.String()))
	if err != nil {
		return 0, err
	}
	// column key -> column value
	for _, c := range t.Cols() {
		colKey := t.RecordKey(recordID, c)
		data, err := t.EncodeValue(r[c.Offset])
		if err != nil {
			return 0, err
		}
		err = txn.Set([]byte(colKey), data)
		if err != nil {
			return 0, err
		}
	}
	variable.GetSessionVars(ctx).AddAffectedRows(1)
	return recordID, nil
}

// Do implements plan.Plan Do interface.
func (r *selectIndexDefaultPlan) Do(ctx context.Context, f plan.RowIterFunc) (err error) {
	var x kv.Index
	switch ix := r.x.(type) {
	case *column.IndexedCol:
		x = ix.X
	default:
		panic("should never happen")
	}

	txn, err := ctx.GetTxn(false)
	if err != nil {
		return err
	}
	en, err := x.SeekFirst(txn)
	if err != nil {
		return types.EOFAsNil(err)
	}
	defer en.Close()

	var id int64
	for {
		k, _, err := en.Next()
		if err != nil {
			return types.EOFAsNil(err)
		}

		id++
		if more, err := f(id, k); !more || err != nil {
			return err
		}
	}
}

func (s *testColumnSuite) checkWriteOnlyColumn(c *C, ctx context.Context, d *ddl, tblInfo *model.TableInfo, handle int64, col *column.Col, row []types.Datum, columnValue interface{}, isDropped bool) {
	t := testGetTable(c, d, s.dbInfo.ID, tblInfo.ID)

	_, err := ctx.GetTxn(true)
	c.Assert(err, IsNil)

	i := int64(0)
	err = t.IterRecords(ctx, t.FirstKey(), t.Cols(), func(h int64, data []types.Datum, cols []*column.Col) (bool, error) {
		c.Assert(data, DeepEquals, row)
		i++
		return true, nil
	})
	c.Assert(err, IsNil)
	c.Assert(i, Equals, int64(1))

	s.checkColumnKVExist(c, ctx, t, handle, col, columnValue, isDropped)

	// Test add a new row.
	_, err = ctx.GetTxn(true)
	c.Assert(err, IsNil)

	newRow := types.MakeDatums(int64(11), int64(22), int64(33))
	handle, err = t.AddRecord(ctx, newRow)
	c.Assert(err, IsNil)

	_, err = ctx.GetTxn(true)
	c.Assert(err, IsNil)

	rows := [][]types.Datum{row, newRow}

	i = int64(0)
	t.IterRecords(ctx, t.FirstKey(), t.Cols(), func(h int64, data []types.Datum, cols []*column.Col) (bool, error) {
		c.Assert(data, DeepEquals, rows[i])
		i++
		return true, nil
	})
	c.Assert(i, Equals, int64(2))

	s.checkColumnKVExist(c, ctx, t, handle, col, columnValue, true)

	// Test remove a row.
	_, err = ctx.GetTxn(true)
	c.Assert(err, IsNil)

	err = t.RemoveRecord(ctx, handle, newRow)
	c.Assert(err, IsNil)

	_, err = ctx.GetTxn(true)
	c.Assert(err, IsNil)

	i = int64(0)
	t.IterRecords(ctx, t.FirstKey(), t.Cols(), func(h int64, data []types.Datum, cols []*column.Col) (bool, error) {
		i++
		return true, nil
	})
	c.Assert(i, Equals, int64(1))

	s.checkColumnKVExist(c, ctx, t, handle, col, columnValue, false)
	s.testGetColumn(c, t, col.Name.L, false)
}

func (d *ddl) deleteTableData(ctx context.Context, t table.Table) error {
	// Remove data
	err := t.Truncate(ctx)
	if err != nil {
		return errors.Trace(err)
	}
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return errors.Trace(err)
	}
	// Remove indices
	for _, v := range t.Indices() {
		if v != nil && v.X != nil {
			if err = v.X.Drop(txn); err != nil {
				return errors.Trace(err)
			}
		}
	}
	// Remove auto ID key
	err = txn.Delete([]byte(meta.AutoIDKey(t.TableID())))

	// Auto ID meta is created when the first time used, so it may not exist.
	if errors2.ErrorEqual(err, kv.ErrNotExist) {
		return nil
	}
	return errors.Trace(err)
}

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

func checkHistoryJobArgs(c *C, ctx context.Context, id int64, args *historyJobArgs) {
	txn, err := ctx.GetTxn(true)
	c.Assert(err, IsNil)
	t := meta.NewMeta(txn)
	historyJob, err := t.GetHistoryDDLJob(id)
	c.Assert(err, IsNil)

	var v int64
	var ids []int64
	tbl := &model.TableInfo{}
	if args.tbl != nil {
		historyJob.DecodeArgs(&v, &tbl)
		c.Assert(v, Equals, args.ver)
		checkEqualTable(c, tbl, args.tbl)
		return
	}
	// only for create schema job
	db := &model.DBInfo{}
	if args.db != nil && len(args.tblIDs) == 0 {
		historyJob.DecodeArgs(&v, &db)
		c.Assert(v, Equals, args.ver)
		c.Assert(db, DeepEquals, args.db)
		return
	}
	// only for drop schema job
	historyJob.DecodeArgs(&v, &db, &ids)
	c.Assert(v, Equals, args.ver)
	c.Assert(db, DeepEquals, args.db)
	for _, id := range ids {
		c.Assert(args.tblIDs, HasKey, id)
		delete(args.tblIDs, id)
	}
	c.Assert(len(args.tblIDs), Equals, 0)
}

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

func (s *testColumnSuite) checkColumnKVExist(ctx context.Context, t table.Table, handle int64, col *table.Column, columnValue interface{}, isExist bool) error {
	txn, err := ctx.GetTxn(true)
	if err != nil {
		return errors.Trace(err)
	}
	defer ctx.CommitTxn()
	key := t.RecordKey(handle)
	data, err := txn.Get(key)
	if !isExist {
		if terror.ErrorEqual(err, kv.ErrNotExist) {
			return nil
		}
	}
	if err != nil {
		return errors.Trace(err)
	}
	colMap := make(map[int64]*types.FieldType)
	colMap[col.ID] = &col.FieldType
	rowMap, err := tablecodec.DecodeRow(data, colMap)
	if err != nil {
		return errors.Trace(err)
	}
	val, ok := rowMap[col.ID]
	if isExist {
		if !ok || val.GetValue() != columnValue {
			return errors.Errorf("%v is not equal to %v", val.GetValue(), columnValue)
		}
	} else {
		if ok {
			return errors.Errorf("column value should not exists")
		}
	}
	return nil
}

func (t *Table) removeRowData(ctx context.Context, h int64) error {
	if err := t.LockRow(ctx, h); err != nil {
		return errors.Trace(err)
	}
	txn, err := ctx.GetTxn(false)
	if err != nil {
		return errors.Trace(err)
	}
	// Remove row's colume one by one
	for _, col := range t.Columns {
		k := t.RecordKey(h, col)
		err = txn.Delete([]byte(k))
		if err != nil {
			if col.State != model.StatePublic && terror.ErrorEqual(err, kv.ErrNotExist) {
				// If the column is not in public state, we may have not added the column,
				// or already deleted the column, so skip ErrNotExist error.
				continue
			}

			return errors.Trace(err)
		}
	}
	// Remove row lock
	err = txn.Delete([]byte(t.RecordKey(h, nil)))
	if err != nil {
		return errors.Trace(err)
	}
	return nil
}

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
}

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

	// use the length of t.Cols() for alignment
	v := make([]interface{}, len(t.Cols()))
	for _, col := range cols {
		if col.State != model.StatePublic {
			return nil, errors.Errorf("Cannot use none public column - %v", cols)
		}

		k := t.RecordKey(h, col)
		data, err := txn.Get([]byte(k))
		if err != nil {
			return nil, errors.Trace(err)
		}

		val, err := t.DecodeValue(data, col)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v[col.Offset] = val
	}
	return v, nil
}