extension: add more informations to extension context

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