extension: add more informations to extension context

br/pkg/task/BUILD.bazel

@@ -96,6 +96,7 @@ go_test(
     flaky = True,
     deps = [
         "//br/pkg/conn",
+        "//br/pkg/errors",
         "//br/pkg/metautil",
         "//br/pkg/restore",
         "//br/pkg/storage",

expression/extension.go

@@ -15,6 +15,7 @@
 package expression
 
 import (
+	"context"
 	"strings"
 	"sync"
 
@@ -104,7 +105,7 @@ func (c *extensionFuncClass) getFunction(ctx sessionctx.Context, args []Expressi
 		return nil, err
 	}
 	bf.tp.SetFlen(c.flen)
-	sig := &extensionFuncSig{bf, c.funcDef}
+	sig := &extensionFuncSig{context.TODO(), bf, c.funcDef}
 	return sig, nil
 }
 
@@ -137,6 +138,7 @@ func (c *extensionFuncClass) checkPrivileges(ctx sessionctx.Context) error {
 var _ extension.FunctionContext = &extensionFuncSig{}
 
 type extensionFuncSig struct {
+	context.Context
 	baseBuiltinFunc
 	extension.FunctionDef
 }

extension/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
         "function.go",
         "manifest.go",
         "registry.go",
+        "session.go",
         "util.go",
     ],
     importpath = "github.com/pingcap/tidb/extension",
@@ -15,7 +16,9 @@ go_library(
         "//sessionctx/variable",
         "//types",
         "//util/chunk",
+        "@com_github_ngaut_pools//:pools",
         "@com_github_pingcap_errors//:errors",
+        "@io_etcd_go_etcd_client_v3//:client",
     ],
 )
 

extension/extensionimpl/BUILD.bazel

@@ -12,5 +12,6 @@ go_library(
         "//util/chunk",
         "//util/sqlexec",
         "@com_github_pingcap_errors//:errors",
+        "@io_etcd_go_etcd_client_v3//:client",
     ],
 )

extension/extensionimpl/bootstrap.go

@@ -23,11 +23,15 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/sqlexec"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 type bootstrapContext struct {
 	context.Context
+
 	sqlExecutor sqlexec.SQLExecutor
+	etcdCli     *clientv3.Client
+	sessionPool extension.SessionPool
 }
 
 func (c *bootstrapContext) ExecuteSQL(ctx context.Context, sql string) (rows []chunk.Row, err error) {
@@ -51,6 +55,14 @@ func (c *bootstrapContext) ExecuteSQL(ctx context.Context, sql string) (rows []c
 	return sqlexec.DrainRecordSet(ctx, rs, 8)
 }
 
+func (c *bootstrapContext) EtcdClient() *clientv3.Client {
+	return c.etcdCli
+}
+
+func (c *bootstrapContext) SessionPool() extension.SessionPool {
+	return c.sessionPool
+}
+
 // Bootstrap bootstrap all extensions
 func Bootstrap(ctx context.Context, do *domain.Domain) error {
 	extensions, err := extension.GetExtensions()
@@ -74,5 +86,10 @@ func Bootstrap(ctx context.Context, do *domain.Domain) error {
 		return errors.Errorf("type '%T' cannot be casted to 'sqlexec.SQLExecutor'", sctx)
 	}
 
-	return extensions.Bootstrap(&bootstrapContext{ctx, executor})
+	return extensions.Bootstrap(&bootstrapContext{
+		Context:     ctx,
+		sessionPool: pool,
+		sqlExecutor: executor,
+		etcdCli:     do.GetEtcdClient(),
+	})
 }

extension/function.go

@@ -15,12 +15,15 @@
 package extension
 
 import (
+	"context"
+
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
 )
 
 // FunctionContext is a interface to provide context to the custom function
 type FunctionContext interface {
+	context.Context
 	EvalArgs(row chunk.Row) ([]types.Datum, error)
 }
 

extension/manifest.go

@@ -17,11 +17,19 @@ package extension
 import (
 	"context"
 
+	"github.com/ngaut/pools"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util/chunk"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
+// SessionPool is the pool for session
+type SessionPool interface {
+	Get() (pools.Resource, error)
+	Put(pools.Resource)
+}
+
 // Option represents an option to initialize an extension
 type Option func(m *Manifest)
 
@@ -66,6 +74,10 @@ type BootstrapContext interface {
 	context.Context
 	// ExecuteSQL is used to execute a sql
 	ExecuteSQL(ctx context.Context, sql string) ([]chunk.Row, error)
+	// EtcdClient returns the etcd client
+	EtcdClient() *clientv3.Client
+	// SessionPool returns the session pool of domain
+	SessionPool() SessionPool
 }
 
 // WithBootstrap specifies the bootstrap func of an extension

server/BUILD.bazel

@@ -32,6 +32,7 @@ go_library(
         "//errno",
         "//executor",
         "//expression",
+        "//extension",
         "//infoschema",
         "//kv",
         "//meta",
@@ -151,6 +152,7 @@ go_test(
         "//errno",
         "//executor",
         "//expression",
+        "//extension",
         "//infoschema",
         "//kv",
         "//meta",

session/BUILD.bazel

@@ -25,6 +25,7 @@ go_library(
         "//errno",
         "//executor",
         "//expression",
+        "//extension",
         "//extension/extensionimpl",
         "//infoschema",
         "//kv",

sessionctx/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
     importpath = "github.com/pingcap/tidb/sessionctx",
     visibility = ["//visibility:public"],
     deps = [
+        "//extension",
         "//kv",
         "//metrics",
         "//parser/model",

testkit/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "//kv",
         "//parser/ast",
         "//parser/terror",
+        "//planner/core",
         "//session",
         "//session/txninfo",
         "//sessionctx/variable",

util/mock/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
     importpath = "github.com/pingcap/tidb/util/mock",
     visibility = ["//visibility:public"],
     deps = [
+        "//extension",
         "//kv",
         "//parser/ast",
         "//parser/model",