expression: introduce optional properties for EvalContext (#51487)

close #51477
pingcap · Mar 11, 2024 · 5f3fc33 · 5f3fc33
1 parent f2ae698
commit 5f3fc33
Show file tree

Hide file tree

Showing 17 changed files with 611 additions and 32 deletions.
diff --git a/pkg/expression/BUILD.bazel b/pkg/expression/BUILD.bazel
@@ -75,6 +75,7 @@ go_library(
         "//pkg/errctx",
         "//pkg/errno",
         "//pkg/expression/context",
+        "//pkg/expression/contextopt",
         "//pkg/extension",
         "//pkg/kv",
         "//pkg/parser",

diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go
@@ -33,6 +33,7 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/expression/contextopt"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/charset"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
@@ -63,6 +64,10 @@ func (b *baseBuiltinFunc) PbCode() tipb.ScalarFuncSig {
 	return b.pbCode
 }
 
+func (*baseBuiltinFunc) RequiredOptionalEvalProps() (set OptionalEvalPropKeySet) {
+	return
+}
+
 // metadata returns the metadata of a function.
 // metadata means some functions contain extra inner fields which will not
 // contain in `tipb.Expr.children` but must be pushed down to coprocessor
@@ -475,6 +480,7 @@ type vecBuiltinFunc interface {
 
 // builtinFunc stands for a particular function signature.
 type builtinFunc interface {
+	contextopt.RequireOptionalEvalProps
 	vecBuiltinFunc
 
 	// evalInt evaluates int result of builtinFunc by given row.

diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go
@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/expression/contextopt"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/model"
@@ -173,12 +174,13 @@ func (c *currentUserFunctionClass) getFunction(ctx BuildContext, args []Expressi
 		return nil, err
 	}
 	bf.tp.SetFlen(64)
-	sig := &builtinCurrentUserSig{bf}
+	sig := &builtinCurrentUserSig{baseBuiltinFunc: bf}
 	return sig, nil
 }
 
 type builtinCurrentUserSig struct {
 	baseBuiltinFunc
+	contextopt.CurrentUserPropReader
 }
 
 func (b *builtinCurrentUserSig) Clone() builtinFunc {
@@ -187,14 +189,22 @@ func (b *builtinCurrentUserSig) Clone() builtinFunc {
 	return newSig
 }
 
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinCurrentUserSig) RequiredOptionalEvalProps() (set OptionalEvalPropKeySet) {
+	return b.CurrentUserPropReader.RequiredOptionalEvalProps()
+}
+
 // evalString evals a builtinCurrentUserSig.
 // See https://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_current-user
-func (b *builtinCurrentUserSig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
-	data := ctx.GetSessionVars()
-	if data == nil || data.User == nil {
+func (b *builtinCurrentUserSig) evalString(ctx EvalContext, _ chunk.Row) (string, bool, error) {
+	user, err := b.CurrentUser(ctx)
+	if err != nil {
+		return "", true, err
+	}
+	if user == nil {
 		return "", true, errors.Errorf("Missing session variable when eval builtin")
 	}
-	return data.User.String(), false, nil
+	return user.String(), false, nil
 }
 
 type currentRoleFunctionClass struct {
@@ -210,12 +220,13 @@ func (c *currentRoleFunctionClass) getFunction(ctx BuildContext, args []Expressi
 		return nil, err
 	}
 	bf.tp.SetFlen(64)
-	sig := &builtinCurrentRoleSig{bf}
+	sig := &builtinCurrentRoleSig{baseBuiltinFunc: bf}
 	return sig, nil
 }
 
 type builtinCurrentRoleSig struct {
 	baseBuiltinFunc
+	contextopt.CurrentUserPropReader
 }
 
 func (b *builtinCurrentRoleSig) Clone() builtinFunc {
@@ -224,24 +235,32 @@ func (b *builtinCurrentRoleSig) Clone() builtinFunc {
 	return newSig
 }
 
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinCurrentRoleSig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.CurrentUserPropReader.RequiredOptionalEvalProps()
+}
+
 // evalString evals a builtinCurrentUserSig.
 // See https://dev.mysql.com/doc/refman/8.0/en/information-functions.html#function_current-role
 func (b *builtinCurrentRoleSig) evalString(ctx EvalContext, row chunk.Row) (res string, isNull bool, err error) {
-	data := ctx.GetSessionVars()
-	if data == nil || data.ActiveRoles == nil {
+	roles, err := b.ActiveRoles(ctx)
+	if err != nil {
+		return "", true, err
+	}
+	if roles == nil {
 		return "", true, errors.Errorf("Missing session variable when eval builtin")
 	}
-	if len(data.ActiveRoles) == 0 {
+	if len(roles) == 0 {
 		return "NONE", false, nil
 	}
 	sortedRes := make([]string, 0, 10)
-	for _, r := range data.ActiveRoles {
+	for _, r := range roles {
 		sortedRes = append(sortedRes, r.String())
 	}
 	slices.Sort(sortedRes)
 	for i, r := range sortedRes {
 		res += r
-		if i != len(data.ActiveRoles)-1 {
+		if i != len(roles)-1 {
 			res += ","
 		}
 	}
@@ -309,12 +328,13 @@ func (c *userFunctionClass) getFunction(ctx BuildContext, args []Expression) (bu
 		return nil, err
 	}
 	bf.tp.SetFlen(64)
-	sig := &builtinUserSig{bf}
+	sig := &builtinUserSig{baseBuiltinFunc: bf}
 	return sig, nil
 }
 
 type builtinUserSig struct {
 	baseBuiltinFunc
+	contextopt.CurrentUserPropReader
 }
 
 func (b *builtinUserSig) Clone() builtinFunc {
@@ -323,14 +343,22 @@ func (b *builtinUserSig) Clone() builtinFunc {
 	return newSig
 }
 
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinUserSig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.CurrentUserPropReader.RequiredOptionalEvalProps()
+}
+
 // evalString evals a builtinUserSig.
 // See https://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_user
-func (b *builtinUserSig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
-	data := ctx.GetSessionVars()
-	if data == nil || data.User == nil {
+func (b *builtinUserSig) evalString(ctx EvalContext, _ chunk.Row) (string, bool, error) {
+	user, err := b.CurrentUser(ctx)
+	if err != nil {
+		return "", true, err
+	}
+	if user == nil {
 		return "", true, errors.Errorf("Missing session variable when eval builtin")
 	}
-	return data.User.LoginString(), false, nil
+	return user.LoginString(), false, nil
 }
 
 type connectionIDFunctionClass struct {

diff --git a/pkg/expression/builtin_info_vec.go b/pkg/expression/builtin_info_vec.go
@@ -106,14 +106,17 @@ func (b *builtinCurrentUserSig) vectorized() bool {
 // See https://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_current-user
 func (b *builtinCurrentUserSig) vecEvalString(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
 	n := input.NumRows()
+	user, err := b.CurrentUser(ctx)
+	if err != nil {
+		return err
+	}
 
-	data := ctx.GetSessionVars()
 	result.ReserveString(n)
-	if data == nil || data.User == nil {
+	if user == nil {
 		return errors.Errorf("Missing session variable when eval builtin")
 	}
 	for i := 0; i < n; i++ {
-		result.AppendString(data.User.String())
+		result.AppendString(user.String())
 	}
 	return nil
 }
@@ -145,21 +148,25 @@ func (b *builtinCurrentRoleSig) vectorized() bool {
 func (b *builtinCurrentRoleSig) vecEvalString(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
 	n := input.NumRows()
 
-	data := ctx.GetSessionVars()
-	if data == nil || data.ActiveRoles == nil {
+	roles, err := b.ActiveRoles(ctx)
+	if err != nil {
+		return err
+	}
+
+	if roles == nil {
 		return errors.Errorf("Missing session variable when eval builtin")
 	}
 
 	result.ReserveString(n)
-	if len(data.ActiveRoles) == 0 {
+	if len(roles) == 0 {
 		for i := 0; i < n; i++ {
 			result.AppendString("NONE")
 		}
 		return nil
 	}
 
 	sortedRes := make([]string, 0, 10)
-	for _, r := range data.ActiveRoles {
+	for _, r := range roles {
 		sortedRes = append(sortedRes, r.String())
 	}
 	slices.Sort(sortedRes)
@@ -178,14 +185,17 @@ func (b *builtinUserSig) vectorized() bool {
 // See https://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_user
 func (b *builtinUserSig) vecEvalString(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
 	n := input.NumRows()
-	data := ctx.GetSessionVars()
-	if data == nil || data.User == nil {
+	user, err := b.CurrentUser(ctx)
+	if err != nil {
+		return err
+	}
+	if user == nil {
 		return errors.Errorf("Missing session variable when eval builtin")
 	}
 
 	result.ReserveString(n)
 	for i := 0; i < n; i++ {
-		result.AppendString(data.User.LoginString())
+		result.AppendString(user.LoginString())
 	}
 	return nil
 }

diff --git a/pkg/expression/context.go b/pkg/expression/context.go
@@ -31,6 +31,18 @@ type EvalContext = context.EvalContext
 // BuildContext is used to build an expression
 type BuildContext = context.BuildContext
 
+// OptionalEvalPropKey is an alias of context.OptionalEvalPropKey
+type OptionalEvalPropKey = context.OptionalEvalPropKey
+
+// OptionalEvalPropProvider is an alias of context.OptionalEvalPropProvider
+type OptionalEvalPropProvider = context.OptionalEvalPropProvider
+
+// OptionalEvalPropKeySet is an alias of context.OptionalEvalPropKeySet
+type OptionalEvalPropKeySet = context.OptionalEvalPropKeySet
+
+// OptionalEvalPropDesc is an alias of context.OptionalEvalPropDesc
+type OptionalEvalPropDesc = context.OptionalEvalPropDesc
+
 func sqlMode(ctx EvalContext) mysql.SQLMode {
 	return ctx.SQLMode()
 }

diff --git a/pkg/expression/context/BUILD.bazel b/pkg/expression/context/BUILD.bazel
@@ -1,8 +1,11 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "context",
-    srcs = ["context.go"],
+    srcs = [
+        "context.go",
+        "optional.go",
+    ],
     importpath = "github.com/pingcap/tidb/pkg/expression/context",
     visibility = ["//visibility:public"],
     deps = [
@@ -15,3 +18,13 @@ go_library(
         "//pkg/types",
     ],
 )
+
+go_test(
+    name = "context_test",
+    timeout = "short",
+    srcs = ["optional_test.go"],
+    embed = [":context"],
+    flaky = True,
+    shard_count = 4,
+    deps = ["@com_github_stretchr_testify//require"],
+)
diff --git a/pkg/expression/context/context.go b/pkg/expression/context/context.go
@@ -71,6 +71,8 @@ type EvalContext interface {
 	GetInfoSchema() infoschema.InfoSchemaMetaVersion
 	// GetDomainInfoSchema returns the latest information schema in domain
 	GetDomainInfoSchema() infoschema.InfoSchemaMetaVersion
+	// GetOptionalPropProvider gets the optional property provider by key
+	GetOptionalPropProvider(OptionalEvalPropKey) (OptionalEvalPropProvider, bool)
 }
 
 // BuildContext is used to build an expression