Golang Context.SetValue方法代码示例

// ExecRestrictedSQL implements SQLHelper interface.
// This is used for executing some restricted sql statements.
func (s *session) ExecRestrictedSQL(ctx context.Context, sql string) (rset.Recordset, error) {
	if ctx.Value(&sqlexec.RestrictedSQLExecutorKeyType{}) != nil {
		// We do not support run this function concurrently.
		// TODO: Maybe we should remove this restriction latter.
		return nil, errors.New("Should not call ExecRestrictedSQL concurrently.")
	}
	statements, err := Compile(ctx, sql)
	if err != nil {
		log.Errorf("Compile %s with error: %v", sql, err)
		return nil, errors.Trace(err)
	}
	if len(statements) != 1 {
		log.Errorf("ExecRestrictedSQL only executes one statement. Too many/few statement in %s", sql)
		return nil, errors.New("Wrong number of statement.")
	}
	st := statements[0]
	// Check statement for some restriction
	// For example only support DML on system meta table.
	// TODO: Add more restrictions.
	log.Debugf("Executing %s [%s]", st.OriginText(), sql)
	ctx.SetValue(&sqlexec.RestrictedSQLExecutorKeyType{}, true)
	defer ctx.ClearValue(&sqlexec.RestrictedSQLExecutorKeyType{})
	rs, err := st.Exec(ctx)
	return rs, errors.Trace(err)
}

// BindSessionVars creates a session vars object and binds it to context.
func BindSessionVars(ctx context.Context) {
	v := &SessionVars{
		Users:         make(map[string]string),
		Systems:       make(map[string]string),
		PreparedStmts: make(map[string]interface{}),
	}

	ctx.SetValue(sessionVarsKey, v)
}

// GetPrewriteValue gets binlog prewrite value in the context.
func GetPrewriteValue(ctx context.Context, createIfNotExists bool) *binlog.PrewriteValue {
	v, ok := ctx.Value(binlogKey).(*binlog.PrewriteValue)
	if !ok && createIfNotExists {
		schemaVer := GetSchemaVersion(ctx)
		v = &binlog.PrewriteValue{SchemaVersion: schemaVer}
		ctx.SetValue(binlogKey, v)
	}
	return v
}

// BindSessionVars creates a session vars object and binds it to context.
func BindSessionVars(ctx context.Context) {
	v := &SessionVars{
		Users:                make(map[string]string),
		systems:              make(map[string]string),
		PreparedStmts:        make(map[uint32]interface{}),
		PreparedStmtNameToID: make(map[string]uint32),
		RetryInfo:            &RetryInfo{},
		StrictSQLMode:        true,
	}
	ctx.SetValue(sessionVarsKey, v)
}

func getDirtyDB(ctx context.Context) *dirtyDB {
	var udb *dirtyDB
	x := ctx.Value(DirtyDBKey)
	if x == nil {
		udb = &dirtyDB{tables: make(map[int64]*dirtyTable)}
		ctx.SetValue(DirtyDBKey, udb)
	} else {
		udb = x.(*dirtyDB)
	}
	return udb
}

func (sq *SubQuery) push(ctx context.Context) {
	var st []*SubQuery
	v := ctx.Value(subQueryStackKey)
	if v == nil {
		st = []*SubQuery{}
	} else {
		// must ok
		st = v.([]*SubQuery)
	}

	st = append(st, sq)
	ctx.SetValue(subQueryStackKey, st)
}

func pushRowStack(ctx context.Context, outDataFields []*field.ResultField, fromDataFields []*field.ResultField) {
	s := getRowStack(ctx)
	if s == nil {
		s = &RowStack{
			items: make([]*rowStackItem, 0, 1),
		}
	}

	s.items = append(s.items, &rowStackItem{
		OutDataFields:  outDataFields,
		FromDataFields: fromDataFields,
	})

	ctx.SetValue(rowStackKey, s)
}

func popRowStack(ctx context.Context) error {
	s := getRowStack(ctx)

	if s == nil || len(s.items) == 0 {
		return errors.Errorf("pop empty row stack")
	}

	n := len(s.items) - 1
	s.items[n] = nil
	s.items = s.items[0:n]

	if len(s.items) == 0 {
		ctx.ClearValue(rowStackKey)
		return nil
	}

	ctx.SetValue(rowStackKey, s)
	return nil
}

func (sq *SubQuery) pop(ctx context.Context) error {
	v := ctx.Value(subQueryStackKey)
	if v == nil {
		return errors.Errorf("pop empty sub query stack")
	}

	st := v.([]*SubQuery)

	// can not empty
	n := len(st) - 1
	if st[n] != sq {
		return errors.Errorf("pop invalid top sub query in stack, want %v, but top is %v", sq, st[n])
	}

	st[n] = nil
	st = st[0:n]
	if len(st) == 0 {
		ctx.ClearValue(subQueryStackKey)
		return nil
	}

	ctx.SetValue(subQueryStackKey, st)
	return nil
}

// BindGlobalVarAccessor binds global var accessor to context.
func BindGlobalVarAccessor(ctx context.Context, accessor GlobalVarAccessor) {
	ctx.SetValue(accessorKey, accessor)
}

// BindExecArgs binds executive args to context.
func BindExecArgs(ctx context.Context, args []interface{}) {
	ctx.SetValue(execArgsKey, args)
}

// SetForUpdate set "select for update" flag.
func SetForUpdate(ctx context.Context) {
	ctx.SetValue(ForUpdateKey, true)
}

// BindAutocommitChecker binds autocommit checker to context.
func BindAutocommitChecker(ctx context.Context, checker Checker) {
	ctx.SetValue(key, checker)
}

// BindPrivilegeChecker binds Checker to context.
func BindPrivilegeChecker(ctx context.Context, pc Checker) {
	ctx.SetValue(key, pc)
}

// SetSchemaVersion sets schema version to the context.
func SetSchemaVersion(ctx context.Context, version int64) {
	ctx.SetValue(schemaVersionKey, version)
}