expression: introduce SessionEvalContext to implement EvalContext (#52091)

close #52089

Author: lcwangchao

pkg/executor/BUILD.bazel

@@ -133,7 +133,6 @@ go_library(
         "//pkg/expression",
         "//pkg/expression/aggregation",
         "//pkg/expression/context",
-        "//pkg/expression/contextimpl",
         "//pkg/infoschema",
         "//pkg/keyspace",
         "//pkg/kv",
@@ -397,7 +396,6 @@ go_test(
         "//pkg/executor/sortexec",
         "//pkg/expression",
         "//pkg/expression/aggregation",
-        "//pkg/expression/contextimpl",
         "//pkg/infoschema",
         "//pkg/kv",
         "//pkg/meta",

pkg/executor/cluster_table_test.go

@@ -27,7 +27,6 @@ import (
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/expression"
-	"github.com/pingcap/tidb/pkg/expression/contextimpl"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/auth"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
@@ -312,9 +311,7 @@ func TestSQLDigestTextRetriever(t *testing.T) {
 		},
 	}
 
-	sqlExec, err := contextimpl.NewSQLExecutor(tk.Session())
-	require.NoError(t, err)
-	err = r.RetrieveLocal(context.Background(), sqlExec)
+	err := r.RetrieveLocal(context.Background(), tk.Session().GetRestrictedSQLExecutor())
 	require.NoError(t, err)
 	require.Equal(t, insertNormalized, r.SQLDigestsMap[insertDigest.String()])
 	require.Equal(t, "", r.SQLDigestsMap[updateDigest.String()])

pkg/executor/infoschema_reader.go

@@ -40,7 +40,6 @@ import (
 	"github.com/pingcap/tidb/pkg/executor/internal/exec"
 	"github.com/pingcap/tidb/pkg/executor/internal/pdhelper"
 	"github.com/pingcap/tidb/pkg/expression"
-	"github.com/pingcap/tidb/pkg/expression/contextimpl"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
@@ -2718,11 +2717,9 @@ func (e *tidbTrxTableRetriever) retrieve(ctx context.Context, sctx sessionctx.Co
 		e.batchRetrieverHelper.batchSize = 1024
 	}
 
-	sqlExec, err := contextimpl.NewSQLExecutor(sctx)
-	if err != nil {
-		return nil, err
-	}
+	sqlExec := sctx.GetRestrictedSQLExecutor()
 
+	var err error
 	// The current TiDB node's address is needed by the CLUSTER_TIDB_TRX table.
 	var instanceAddr string
 	if e.table.Name.O == infoschema.ClusterTableTiDBTrx {
@@ -2871,12 +2868,7 @@ func (r *dataLockWaitsTableRetriever) retrieve(ctx context.Context, sctx session
 				}
 			}
 
-			sqlExec, err := contextimpl.NewSQLExecutor(sctx)
-			if err != nil {
-				return errors.Trace(err)
-			}
-
-			err = sqlRetriever.RetrieveGlobal(ctx, sqlExec)
+			err := sqlRetriever.RetrieveGlobal(ctx, sctx.GetRestrictedSQLExecutor())
 			if err != nil {
 				return errors.Trace(err)
 			}
@@ -3074,11 +3066,7 @@ func (r *deadlocksTableRetriever) retrieve(ctx context.Context, sctx sessionctx.
 		}
 		// Retrieve the SQL texts if necessary.
 		if sqlRetriever != nil {
-			sqlExec, err := contextimpl.NewSQLExecutor(sctx)
-			if err != nil {
-				return errors.Trace(err)
-			}
-			err1 := sqlRetriever.RetrieveGlobal(ctx, sqlExec)
+			err1 := sqlRetriever.RetrieveGlobal(ctx, sctx.GetRestrictedSQLExecutor())
 			if err1 != nil {
 				return errors.Trace(err1)
 			}

pkg/expression/contextimpl/BUILD.bazel

@@ -10,7 +10,6 @@ go_library(
         "//pkg/expression/context",
         "//pkg/expression/contextopt",
         "//pkg/infoschema/context",
-        "//pkg/parser/ast",
         "//pkg/parser/auth",
         "//pkg/parser/model",
         "//pkg/parser/mysql",
@@ -20,11 +19,8 @@ go_library(
         "//pkg/sessionctx/variable",
         "//pkg/types",
         "//pkg/util",
-        "//pkg/util/chunk",
         "//pkg/util/intest",
         "//pkg/util/logutil",
-        "//pkg/util/sqlexec",
-        "@com_github_pingcap_errors//:errors",
         "@com_github_tikv_client_go_v2//oracle",
         "@org_uber_go_zap//:zap",
     ],

pkg/expression/contextimpl/sessionctx.go

@@ -19,12 +19,10 @@ import (
 	"math"
 	"time"
 
-	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/errctx"
 	exprctx "github.com/pingcap/tidb/pkg/expression/context"
 	"github.com/pingcap/tidb/pkg/expression/contextopt"
 	infoschema "github.com/pingcap/tidb/pkg/infoschema/context"
-	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/auth"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
@@ -34,10 +32,8 @@ import (
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util"
-	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/pingcap/tidb/pkg/util/logutil"
-	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"github.com/tikv/client-go/v2/oracle"
 	"go.uber.org/zap"
 )
@@ -51,92 +47,104 @@ var _ exprctx.BuildContext = struct {
 
 // ExprCtxExtendedImpl extends the sessionctx.Context to implement `expression.BuildContext`
 type ExprCtxExtendedImpl struct {
-	sctx  sessionctx.Context
-	props contextopt.OptionalEvalPropProviders
+	*SessionEvalContext
 }
 
 // NewExprExtendedImpl creates a new ExprCtxExtendedImpl.
 func NewExprExtendedImpl(sctx sessionctx.Context) *ExprCtxExtendedImpl {
-	impl := &ExprCtxExtendedImpl{sctx: sctx}
+	return &ExprCtxExtendedImpl{
+		SessionEvalContext: NewSessionEvalContext(sctx),
+	}
+}
+
+// SessionEvalContext implements the `expression.EvalContext` interface to provide evaluation context in session.
+type SessionEvalContext struct {
+	sctx  sessionctx.Context
+	props contextopt.OptionalEvalPropProviders
+}
+
+// NewSessionEvalContext creates a new SessionEvalContext.
+func NewSessionEvalContext(sctx sessionctx.Context) *SessionEvalContext {
+	ctx := &SessionEvalContext{sctx: sctx}
 	// set all optional properties
-	impl.setOptionalProp(currentUserProp(sctx))
-	impl.setOptionalProp(contextopt.NewSessionVarsProvider(sctx))
-	impl.setOptionalProp(infoSchemaProp(sctx))
-	impl.setOptionalProp(contextopt.KVStorePropProvider(sctx.GetStore))
-	impl.setOptionalProp(sqlExecutorProp(sctx))
-	impl.setOptionalProp(sequenceOperatorProp(sctx))
-	impl.setOptionalProp(contextopt.NewAdvisoryLockPropProvider(sctx))
-	impl.setOptionalProp(contextopt.DDLOwnerInfoProvider(sctx.IsDDLOwner))
+	ctx.setOptionalProp(currentUserProp(sctx))
+	ctx.setOptionalProp(contextopt.NewSessionVarsProvider(sctx))
+	ctx.setOptionalProp(infoSchemaProp(sctx))
+	ctx.setOptionalProp(contextopt.KVStorePropProvider(sctx.GetStore))
+	ctx.setOptionalProp(sqlExecutorProp(sctx))
+	ctx.setOptionalProp(sequenceOperatorProp(sctx))
+	ctx.setOptionalProp(contextopt.NewAdvisoryLockPropProvider(sctx))
+	ctx.setOptionalProp(contextopt.DDLOwnerInfoProvider(sctx.IsDDLOwner))
 	// When EvalContext is created from a session, it should contain all the optional properties.
-	intest.Assert(impl.props.PropKeySet().IsFull())
-	return impl
+	intest.Assert(ctx.props.PropKeySet().IsFull())
+	return ctx
 }
 
-func (ctx *ExprCtxExtendedImpl) setOptionalProp(prop exprctx.OptionalEvalPropProvider) {
+func (ctx *SessionEvalContext) setOptionalProp(prop exprctx.OptionalEvalPropProvider) {
 	intest.AssertFunc(func() bool {
 		return !ctx.props.Contains(prop.Desc().Key())
 	})
 	ctx.props.Add(prop)
 }
 
 // CtxID returns the context id.
-func (ctx *ExprCtxExtendedImpl) CtxID() uint64 {
+func (ctx *SessionEvalContext) CtxID() uint64 {
 	return ctx.sctx.GetSessionVars().StmtCtx.CtxID()
 }
 
 // SQLMode returns the sql mode
-func (ctx *ExprCtxExtendedImpl) SQLMode() mysql.SQLMode {
+func (ctx *SessionEvalContext) SQLMode() mysql.SQLMode {
 	return ctx.sctx.GetSessionVars().SQLMode
 }
 
 // TypeCtx returns the types.Context
-func (ctx *ExprCtxExtendedImpl) TypeCtx() types.Context {
+func (ctx *SessionEvalContext) TypeCtx() types.Context {
 	return ctx.sctx.GetSessionVars().StmtCtx.TypeCtx()
 }
 
 // ErrCtx returns the errctx.Context
-func (ctx *ExprCtxExtendedImpl) ErrCtx() errctx.Context {
+func (ctx *SessionEvalContext) ErrCtx() errctx.Context {
 	return ctx.sctx.GetSessionVars().StmtCtx.ErrCtx()
 }
 
 // Location returns the timezone info
-func (ctx *ExprCtxExtendedImpl) Location() *time.Location {
+func (ctx *SessionEvalContext) Location() *time.Location {
 	tc := ctx.TypeCtx()
 	return tc.Location()
 }
 
 // AppendWarning append warnings to the context.
-func (ctx *ExprCtxExtendedImpl) AppendWarning(err error) {
+func (ctx *SessionEvalContext) AppendWarning(err error) {
 	ctx.sctx.GetSessionVars().StmtCtx.AppendWarning(err)
 }
 
 // WarningCount gets warning count.
-func (ctx *ExprCtxExtendedImpl) WarningCount() int {
+func (ctx *SessionEvalContext) WarningCount() int {
 	return int(ctx.sctx.GetSessionVars().StmtCtx.WarningCount())
 }
 
 // TruncateWarnings truncates warnings begin from start and returns the truncated warnings.
-func (ctx *ExprCtxExtendedImpl) TruncateWarnings(start int) []stmtctx.SQLWarn {
+func (ctx *SessionEvalContext) TruncateWarnings(start int) []stmtctx.SQLWarn {
 	return ctx.sctx.GetSessionVars().StmtCtx.TruncateWarnings(start)
 }
 
 // CurrentDB returns the current database name
-func (ctx *ExprCtxExtendedImpl) CurrentDB() string {
+func (ctx *SessionEvalContext) CurrentDB() string {
 	return ctx.sctx.GetSessionVars().CurrentDB
 }
 
 // CurrentTime returns the current time
-func (ctx *ExprCtxExtendedImpl) CurrentTime() (time.Time, error) {
+func (ctx *SessionEvalContext) CurrentTime() (time.Time, error) {
 	return getStmtTimestamp(ctx.sctx)
 }
 
 // GetMaxAllowedPacket returns the value of the 'max_allowed_packet' system variable.
-func (ctx *ExprCtxExtendedImpl) GetMaxAllowedPacket() uint64 {
+func (ctx *SessionEvalContext) GetMaxAllowedPacket() uint64 {
 	return ctx.sctx.GetSessionVars().MaxAllowedPacket
 }
 
 // GetDefaultWeekFormatMode returns the value of the 'default_week_format' system variable.
-func (ctx *ExprCtxExtendedImpl) GetDefaultWeekFormatMode() string {
+func (ctx *SessionEvalContext) GetDefaultWeekFormatMode() string {
 	mode, ok := ctx.sctx.GetSessionVars().GetSystemVar(variable.DefaultWeekFormat)
 	if !ok || mode == "" {
 		return "0"
@@ -145,22 +153,22 @@ func (ctx *ExprCtxExtendedImpl) GetDefaultWeekFormatMode() string {
 }
 
 // GetDivPrecisionIncrement returns the specified value of DivPrecisionIncrement.
-func (ctx *ExprCtxExtendedImpl) GetDivPrecisionIncrement() int {
+func (ctx *SessionEvalContext) GetDivPrecisionIncrement() int {
 	return ctx.sctx.GetSessionVars().GetDivPrecisionIncrement()
 }
 
 // GetOptionalPropSet gets the optional property set from context
-func (ctx *ExprCtxExtendedImpl) GetOptionalPropSet() exprctx.OptionalEvalPropKeySet {
+func (ctx *SessionEvalContext) GetOptionalPropSet() exprctx.OptionalEvalPropKeySet {
 	return ctx.props.PropKeySet()
 }
 
 // GetOptionalPropProvider gets the optional property provider by key
-func (ctx *ExprCtxExtendedImpl) GetOptionalPropProvider(key exprctx.OptionalEvalPropKey) (exprctx.OptionalEvalPropProvider, bool) {
+func (ctx *SessionEvalContext) GetOptionalPropProvider(key exprctx.OptionalEvalPropKey) (exprctx.OptionalEvalPropProvider, bool) {
 	return ctx.props.Get(key)
 }
 
 // RequestVerification verifies user privilege
-func (ctx *ExprCtxExtendedImpl) RequestVerification(db, table, column string, priv mysql.PrivilegeType) bool {
+func (ctx *SessionEvalContext) RequestVerification(db, table, column string, priv mysql.PrivilegeType) bool {
 	checker := privilege.GetPrivilegeManager(ctx.sctx)
 	if checker == nil {
 		return true
@@ -169,7 +177,7 @@ func (ctx *ExprCtxExtendedImpl) RequestVerification(db, table, column string, pr
 }
 
 // RequestDynamicVerification verifies user privilege for a DYNAMIC privilege.
-func (ctx *ExprCtxExtendedImpl) RequestDynamicVerification(privName string, grantable bool) bool {
+func (ctx *SessionEvalContext) RequestDynamicVerification(privName string, grantable bool) bool {
 	checker := privilege.GetPrivilegeManager(ctx.sctx)
 	if checker == nil {
 		return true
@@ -223,27 +231,9 @@ func infoSchemaProp(sctx sessionctx.Context) contextopt.InfoSchemaPropProvider {
 	}
 }
 
-type sqlExecutor struct {
-	exec sqlexec.RestrictedSQLExecutor
-}
-
-// NewSQLExecutor creates a new SQLExecutor.
-func NewSQLExecutor(sctx sessionctx.Context) (contextopt.SQLExecutor, error) {
-	if e, ok := sctx.(sqlexec.RestrictedSQLExecutor); ok {
-		return &sqlExecutor{exec: e}, nil
-	}
-	return nil, errors.Errorf("'%T' cannot be casted to sqlexec.RestrictedSQLExecutor", sctx)
-}
-
-func (e *sqlExecutor) ExecRestrictedSQL(
-	ctx context.Context, sql string, args ...any,
-) ([]chunk.Row, []*ast.ResultField, error) {
-	return e.exec.ExecRestrictedSQL(ctx, nil, sql, args...)
-}
-
 func sqlExecutorProp(sctx sessionctx.Context) contextopt.SQLExecutorPropProvider {
 	return func() (contextopt.SQLExecutor, error) {
-		return NewSQLExecutor(sctx)
+		return sctx.GetRestrictedSQLExecutor(), nil
 	}
 }
 

pkg/expression/contextimpl/sessionctx_test.go

@@ -35,7 +35,7 @@ import (
 	"github.com/tikv/client-go/v2/oracle"
 )
 
-func TestEvalContextImplWithSessionCtx(t *testing.T) {
+func TestSessionEvalContextBasic(t *testing.T) {
 	ctx := mock.NewContext()
 	vars := ctx.GetSessionVars()
 	sc := vars.StmtCtx
@@ -93,7 +93,7 @@ func TestEvalContextImplWithSessionCtx(t *testing.T) {
 	require.Equal(t, "err1", warnings[0].Err.Error())
 }
 
-func TestEvalContextImplCurrentTime(t *testing.T) {
+func TestSessionEvalContextCurrentTime(t *testing.T) {
 	ctx := mock.NewContext()
 	vars := ctx.GetSessionVars()
 	sc := vars.StmtCtx
@@ -162,7 +162,7 @@ func (m *mockPrivManager) RequestDynamicVerification(
 	return m.Called(activeRoles, privName, grantable).Bool(0)
 }
 
-func TestEvalContextImplPrivilegeCheck(t *testing.T) {
+func TestSessionEvalContextPrivilegeCheck(t *testing.T) {
 	ctx := mock.NewContext()
 	impl := contextimpl.NewExprExtendedImpl(ctx)
 	activeRoles := []*auth.RoleIdentity{
@@ -212,7 +212,7 @@ func getProvider[T context.OptionalEvalPropProvider](
 	return p
 }
 
-func TestEvalContextImplWithSessionCtxForOptProps(t *testing.T) {
+func TestSessionEvalContextOptProps(t *testing.T) {
 	ctx := mock.NewContext()
 	impl := contextimpl.NewExprExtendedImpl(ctx)
 

pkg/expression/contextopt/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "//pkg/sessionctx/variable",
         "//pkg/util/chunk",
         "//pkg/util/intest",
+        "//pkg/util/sqlexec",
         "@com_github_pingcap_errors//:errors",
     ],
 )

pkg/expression/contextopt/sqlexec.go

@@ -20,12 +20,21 @@ import (
 	exprctx "github.com/pingcap/tidb/pkg/expression/context"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/sqlexec"
 )
 
+// SQLExecutor provides a subset of methods in RestrictedSQLExecutor.
+var _ SQLExecutor = sqlexec.RestrictedSQLExecutor(nil)
+
 // SQLExecutor is the interface for SQL executing in expression.
-// We do not `sqlexec.SQLExecutor` here to avoid to introduce too many dependencies in `sessionctx.Context`
+// We do not `sqlexec.SQLExecutor` to limit expression to use specified methods only.
 type SQLExecutor interface {
-	ExecRestrictedSQL(ctx context.Context, sql string, args ...any) ([]chunk.Row, []*ast.ResultField, error)
+	ExecRestrictedSQL(
+		ctx context.Context,
+		opts []sqlexec.OptionFuncAlias,
+		sql string,
+		args ...any,
+	) ([]chunk.Row, []*ast.ResultField, error)
 }
 
 // SQLExecutorPropProvider provides the SQLExecutor