Golang context.Context类代码示例

// GetDomain gets domain from context.
func GetDomain(ctx context.Context) *domain.Domain {
	v, ok := ctx.Value(domainKey).(*domain.Domain)
	if !ok {
		return nil
	}
	return v
}

// GetGlobalVarAccessor gets accessor from ctx.
func GetGlobalVarAccessor(ctx context.Context) GlobalVarAccessor {
	v, ok := ctx.Value(accessorKey).(GlobalVarAccessor)
	if !ok {
		panic("Miss global sysvar accessor")
	}
	return v
}

// 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 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 *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 runStmt(ctx context.Context, s stmt.Statement, args ...interface{}) (rset.Recordset, error) {
	var err error
	var rs rset.Recordset
	// before every execution, we must clear affectedrows.
	variable.GetSessionVars(ctx).SetAffectedRows(0)
	switch s.(type) {
	case *stmts.PreparedStmt:
		ps := s.(*stmts.PreparedStmt)
		return runPreparedStmt(ctx, ps)
	case *stmts.ExecuteStmt:
		es := s.(*stmts.ExecuteStmt)
		rs, err = runExecute(ctx, es, args...)
		if err != nil {
			return nil, errors.Trace(err)
		}
	default:
		if s.IsDDL() {
			err = ctx.FinishTxn(false)
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		stmt.BindExecArgs(ctx, args)
		rs, err = s.Exec(ctx)
		stmt.ClearExecArgs(ctx)
	}
	// MySQL DDL should be auto-commit
	if err == nil && (s.IsDDL() || variable.ShouldAutocommit(ctx)) {
		err = ctx.FinishTxn(false)
	}
	return rs, errors.Trace(err)
}

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

func (af *aggFunction) streamUpdateSum(row []types.Datum, ectx context.Context) error {
	ctx := af.getStreamedContext()
	a := af.Args[0]
	value, err := a.Eval(row, ectx)
	if err != nil {
		return errors.Trace(err)
	}
	if value.IsNull() {
		return nil
	}
	if af.Distinct {
		d, err1 := ctx.DistinctChecker.Check([]interface{}{value.GetValue()})
		if err1 != nil {
			return errors.Trace(err1)
		}
		if !d {
			return nil
		}
	}
	ctx.Value, err = types.CalculateSum(ectx.GetSessionVars().StmtCtx, ctx.Value, value)
	if err != nil {
		return errors.Trace(err)
	}
	ctx.Count++
	return nil
}

func (af *avgFunction) updateAvg(row []types.Datum, groupKey []byte, ectx context.Context) error {
	ctx := af.getContext(groupKey)
	a := af.Args[1]
	value, err := a.Eval(row, ectx)
	if err != nil {
		return errors.Trace(err)
	}
	if value.IsNull() {
		return nil
	}
	if af.Distinct {
		d, err1 := ctx.DistinctChecker.Check([]interface{}{value.GetValue()})
		if err1 != nil {
			return errors.Trace(err1)
		}
		if !d {
			return nil
		}
	}
	ctx.Value, err = types.CalculateSum(ectx.GetSessionVars().StmtCtx, ctx.Value, value)
	if err != nil {
		return errors.Trace(err)
	}
	count, err := af.Args[0].Eval(row, ectx)
	if err != nil {
		return errors.Trace(err)
	}
	ctx.Count += count.GetInt64()
	return nil
}

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

func (s *testColumnChangeSuite) checkAddPublic(d *ddl, ctx context.Context, writeOnlyTable, publicTable table.Table) error {
	// publicTable Insert t values (4, 4, 4)
	h, err := publicTable.AddRecord(ctx, types.MakeDatums(4, 4, 4))
	if err != nil {
		return errors.Trace(err)
	}
	err = ctx.CommitTxn()
	if err != nil {
		return errors.Trace(err)
	}
	// writeOnlyTable update t set c1 = 3 where c1 = 4
	oldRow, err := writeOnlyTable.RowWithCols(ctx, h, writeOnlyTable.WritableCols())
	if err != nil {
		return errors.Trace(err)
	}
	if len(oldRow) != 3 {
		return errors.Errorf("%v", oldRow)
	}
	newRow := types.MakeDatums(3, 4, oldRow[2].GetValue())
	err = writeOnlyTable.UpdateRecord(ctx, h, oldRow, newRow, touchedMap(writeOnlyTable))
	if err != nil {
		return errors.Trace(err)
	}
	err = ctx.CommitTxn()
	if err != nil {
		return errors.Trace(err)
	}
	// publicTable select * from t, make sure the new c3 value 4 is not overwritten to default value 3.
	err = checkResult(ctx, publicTable, testutil.RowsWithSep(" ", "2 3 3", "3 4 4"))
	if err != nil {
		return errors.Trace(err)
	}
	return nil
}

// GetSchemaVersion gets schema version in the context.
func GetSchemaVersion(ctx context.Context) int64 {
	v, ok := ctx.Value(schemaVersionKey).(int64)
	if !ok {
		log.Error("get schema version failed")
	}
	return v
}

// GetTiDBSystemVar get variable value for name.
// The variable should be a TiDB specific system variable (The vars in tidbSysVars map).
// If the session scope variable is not set, it will get global scope value and fill session scope value.
func (s *SessionVars) GetTiDBSystemVar(ctx context.Context, name string) (string, error) {
	key := strings.ToLower(name)
	_, ok := tidbSysVars[key]
	if !ok {
		return "", errors.Errorf("%s is not a TiDB specific system variable.", name)
	}

	sVal, ok := s.systems[key]
	if ok {
		return sVal, nil
	}

	if ctx.Value(context.Initing) != nil {
		// When running bootstrap or upgrade job, we should not access global storage.
		return SysVars[key].Value, nil
	}

	if key == DistSQLScanConcurrencyVar {
		// Get global variable need to scan table which depends on DistSQLScanConcurrencyVar.
		// So we should add it here to break the dependency loop.
		s.systems[DistSQLScanConcurrencyVar] = SysVars[key].Value
	}

	globalVars := GetGlobalVarAccessor(ctx)
	globalVal, err := globalVars.GetGlobalSysVar(ctx, key)
	if err != nil {
		if key == DistSQLScanConcurrencyVar {
			// Clean up.
			delete(s.systems, DistSQLScanConcurrencyVar)
		}
		return "", errors.Trace(err)
	}
	s.systems[key] = globalVal
	return globalVal, nil
}

// ShouldAutocommit gets checker from ctx and checks if it should autocommit.
func ShouldAutocommit(ctx context.Context) (bool, error) {
	v, ok := ctx.Value(key).(Checker)
	if !ok {
		panic("Miss autocommit checker")
	}
	return v.ShouldAutocommit(ctx)
}

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