sql: implement SHOW BR JOBS <n:int> and CANCEL BR JOB <n:int>

errno/errcode.go

@@ -1062,6 +1062,7 @@ const (
 	ErrLoadDataJobNotFound                 = 8170
 	ErrLoadDataInvalidOperation            = 8171
 	ErrLoadDataLocalUnsupportedOption      = 8172
+	ErrBRJobNotFound                       = 8173
 
 	// Error codes used by TiDB ddl package
 	ErrUnsupportedDDLOperation            = 8200

errno/errname.go

@@ -1095,6 +1095,7 @@ var MySQLErrName = map[uint16]*mysql.ErrMessage{
 	ErrBRIERestoreFailed: mysql.Message("Restore failed: %s", nil),
 	ErrBRIEImportFailed:  mysql.Message("Import failed: %s", nil),
 	ErrBRIEExportFailed:  mysql.Message("Export failed: %s", nil),
+	ErrBRJobNotFound:     mysql.Message("BRIE Job %d not found", nil),
 
 	ErrInvalidTableSample: mysql.Message("Invalid TABLESAMPLE: %s", nil),
 

errors.toml

@@ -1756,6 +1756,11 @@ error = '''
 Unsupported option for LOAD DATA LOCAL INFILE: %s
 '''
 
+["executor:8173"]
+error = '''
+BRIE Job %d not found
+'''
+
 ["executor:8212"]
 error = '''
 Failed to split region ranges: %s

executor/brie.go

@@ -24,8 +24,10 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/encryptionpb"
+	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/br/pkg/glue"
 	"github.com/pingcap/tidb/br/pkg/storage"
 	"github.com/pingcap/tidb/br/pkg/task"
@@ -37,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/format"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/parser/terror"
@@ -52,6 +55,7 @@ import (
 	filter "github.com/pingcap/tidb/util/table-filter"
 	"github.com/tikv/client-go/v2/oracle"
 	pd "github.com/tikv/pd/client"
+	"go.uber.org/zap"
 )
 
 const clearInterval = 10 * time.Minute
@@ -94,11 +98,17 @@ func (p *brieTaskProgress) GetCurrent() int64 {
 // Close implements glue.Progress
 func (p *brieTaskProgress) Close() {
 	p.lock.Lock()
+	current := atomic.LoadInt64(&p.current)
+	if current < p.total {
+		p.cmd = fmt.Sprintf("%s Cacneled", p.cmd)
+	}
 	atomic.StoreInt64(&p.current, p.total)
 	p.lock.Unlock()
 }
 
 type brieTaskInfo struct {
+	id          uint64
+	query       string
 	queueTime   types.Time
 	execTime    types.Time
 	finishTime  types.Time
@@ -149,10 +159,19 @@ func (bq *brieQueue) registerTask(
 
 	taskID := atomic.AddUint64(&bq.nextID, 1)
 	bq.tasks.Store(taskID, item)
+	info.id = taskID
 
 	return taskCtx, taskID
 }
 
+// query task queries a task from the queue.
+func (bq *brieQueue) queryTask(taskID uint64) (*brieTaskInfo, bool) {
+	if item, ok := bq.tasks.Load(taskID); ok {
+		return item.(*brieQueueItem).info, true
+	}
+	return nil, false
+}
+
 // acquireTask prepares to execute a BRIE task. Only one BRIE task can be
 // executed at a time, and this function blocks until the task is ready.
 //
@@ -176,12 +195,16 @@ func (bq *brieQueue) releaseTask() {
 	<-bq.workerCh
 }
 
-func (bq *brieQueue) cancelTask(taskID uint64) {
+func (bq *brieQueue) cancelTask(taskID uint64) bool {
 	item, ok := bq.tasks.Load(taskID)
 	if !ok {
-		return
+		return false
 	}
-	item.(*brieQueueItem).cancel()
+	i := item.(*brieQueueItem)
+	i.cancel()
+	i.progress.Close()
+	log.Info("BRIE job canceled.", zap.Uint64("ID", i.info.id))
+	return true
 }
 
 func (bq *brieQueue) clearTask(sc *stmtctx.StatementContext) {
@@ -223,8 +246,23 @@ func (b *executorBuilder) buildBRIE(s *ast.BRIEStmt, schema *expression.Schema)
 	}
 
 	if s.Kind == ast.BRIEKindShowBackupMeta {
-		e.fillByShowMetadata(s)
-		return e
+		return execOnce(&showMetaExec{
+			showConfig: buildShowMetadataConfigFrom(s),
+		})
+	}
+
+	if s.Kind == ast.BRIEKindShowQuery {
+		return execOnce(&showQueryExec{
+			baseExecutor: newBaseExecutor(b.ctx, schema, 0),
+			targetID:     uint64(s.JobID),
+		})
+	}
+
+	if s.Kind == ast.BRIEKindCancelJob {
+		return &cancelJobExec{
+			baseExecutor: newBaseExecutor(b.ctx, schema, 0),
+			targetID:     uint64(s.JobID),
+		}
 	}
 
 	tidbCfg := config.GetGlobalConfig()
@@ -308,6 +346,11 @@ func (b *executorBuilder) buildBRIE(s *ast.BRIEStmt, schema *expression.Schema)
 	// is expected to be performed insensitive.
 	cfg.TableFilter = filter.CaseInsensitive(cfg.TableFilter)
 
+	// We cannot directly use the query string, or the secret may be print.
+	// NOTE: the ownership of `s.Storage` is taken here.
+	s.Storage = e.info.storage
+	e.info.query = restoreQuery(s)
+
 	switch s.Kind {
 	case ast.BRIEKindBackup:
 		e.backupCfg = &task.BackupConfig{Config: cfg}
@@ -354,6 +397,67 @@ func (b *executorBuilder) buildBRIE(s *ast.BRIEStmt, schema *expression.Schema)
 	return e
 }
 
+// oneshotExecutor warps a executor, making its `Next` would only be called once.
+type oneshotExecutor struct {
+	Executor
+	finished bool
+}
+
+func (o *oneshotExecutor) Next(ctx context.Context, req *chunk.Chunk) error {
+	if o.finished {
+		req.Reset()
+		return nil
+	}
+
+	if err := o.Executor.Next(ctx, req); err != nil {
+		return err
+	}
+	o.finished = true
+	return nil
+}
+
+func execOnce(ex Executor) Executor {
+	return &oneshotExecutor{Executor: ex}
+}
+
+type showQueryExec struct {
+	baseExecutor
+
+	targetID uint64
+}
+
+func (s *showQueryExec) Next(ctx context.Context, req *chunk.Chunk) error {
+	req.Reset()
+
+	tsk, ok := globalBRIEQueue.queryTask(s.targetID)
+	if !ok {
+		return nil
+	}
+
+	req.AppendString(0, tsk.query)
+	return nil
+}
+
+type cancelJobExec struct {
+	baseExecutor
+
+	targetID uint64
+}
+
+func (s cancelJobExec) Next(ctx context.Context, req *chunk.Chunk) error {
+	req.Reset()
+	if !globalBRIEQueue.cancelTask(s.targetID) {
+		s.ctx.GetSessionVars().StmtCtx.AppendWarning(exeerrors.ErrLoadDataJobNotFound.FastGenByArgs(s.targetID))
+	}
+	return nil
+}
+
+type showMetaExec struct {
+	baseExecutor
+
+	showConfig show.Config
+}
+
 // BRIEExec represents an executor for BRIE statements (BACKUP, RESTORE, etc)
 type BRIEExec struct {
 	baseExecutor
@@ -364,23 +468,23 @@ type BRIEExec struct {
 	info       *brieTaskInfo
 }
 
-func (e *BRIEExec) fillByShowMetadata(s *ast.BRIEStmt) {
+func buildShowMetadataConfigFrom(s *ast.BRIEStmt) show.Config {
 	if s.Kind != ast.BRIEKindShowBackupMeta {
 		panic(fmt.Sprintf("precondition failed: `fillByShowMetadata` should always called by a ast.BRIEKindShowBackupMeta, but it is %s.", s.Kind))
 	}
 
 	store := s.Storage
-	e.showConfig = &show.Config{
+	cfg := show.Config{
 		Storage: store,
 		Cipher: backuppb.CipherInfo{
 			CipherType: encryptionpb.EncryptionMethod_PLAINTEXT,
 		},
 	}
-	e.info.kind = ast.BRIEKindShowBackupMeta
+	return cfg
 }
 
-func (e *BRIEExec) runShowMetadata(ctx context.Context, req *chunk.Chunk) error {
-	exe, err := show.CreateExec(ctx, *e.showConfig)
+func (e *showMetaExec) Next(ctx context.Context, req *chunk.Chunk) error {
+	exe, err := show.CreateExec(ctx, e.showConfig)
 	if err != nil {
 		return errors.Annotate(err, "failed to create show exec")
 	}
@@ -404,7 +508,6 @@ func (e *BRIEExec) runShowMetadata(ctx context.Context, req *chunk.Chunk) error
 		}
 		req.AppendTime(5, types.NewTime(types.FromGoTime(endTime), mysql.TypeDatetime, 0))
 	}
-	e.info = nil
 	return nil
 }
 
@@ -415,21 +518,20 @@ func (e *BRIEExec) Next(ctx context.Context, req *chunk.Chunk) error {
 		return nil
 	}
 
-	if e.info.kind == ast.BRIEKindShowBackupMeta {
-		// This should be able to execute without the queue.
-		// NOTE: maybe extract the procedure of executing task in queue
-		// into a function, make it more tidy.
-		return e.runShowMetadata(ctx, req)
-	}
-
 	bq := globalBRIEQueue
 	bq.clearTask(e.ctx.GetSessionVars().StmtCtx)
 
 	e.info.connID = e.ctx.GetSessionVars().ConnectionID
 	e.info.queueTime = types.CurrentTime(mysql.TypeDatetime)
 	taskCtx, taskID := bq.registerTask(ctx, e.info)
 	defer bq.cancelTask(taskID)
-
+	failpoint.Inject("block-on-brie", func() {
+		log.Warn("You shall not pass, nya. :3")
+		<-taskCtx.Done()
+		if taskCtx.Err() != nil {
+			failpoint.Return(taskCtx.Err())
+		}
+	})
 	// manually monitor the Killed status...
 	go func() {
 		ticker := time.NewTicker(3 * time.Second)
@@ -502,17 +604,18 @@ func (e *ShowExec) fetchShowBRIE(kind ast.BRIEKind) error {
 			item.progress.lock.Lock()
 			defer item.progress.lock.Unlock()
 			current := atomic.LoadInt64(&item.progress.current)
-			e.result.AppendString(0, item.info.storage)
-			e.result.AppendString(1, item.progress.cmd)
-			e.result.AppendFloat64(2, 100.0*float64(current)/float64(item.progress.total))
-			e.result.AppendTime(3, item.info.queueTime)
-			e.result.AppendTime(4, item.info.execTime)
-			e.result.AppendTime(5, item.info.finishTime)
-			e.result.AppendUint64(6, item.info.connID)
+			e.result.AppendUint64(0, item.info.id)
+			e.result.AppendString(1, item.info.storage)
+			e.result.AppendString(2, item.progress.cmd)
+			e.result.AppendFloat64(3, 100.0*float64(current)/float64(item.progress.total))
+			e.result.AppendTime(4, item.info.queueTime)
+			e.result.AppendTime(5, item.info.execTime)
+			e.result.AppendTime(6, item.info.finishTime)
+			e.result.AppendUint64(7, item.info.connID)
 			if len(item.info.message) > 0 {
-				e.result.AppendString(7, item.info.message)
+				e.result.AppendString(8, item.info.message)
 			} else {
-				e.result.AppendNull(7)
+				e.result.AppendNull(8)
 			}
 		}
 		return true
@@ -657,3 +760,12 @@ func (gs *tidbGlueSession) UseOneShotSession(store kv.Storage, closeDomain bool,
 	// in SQL backup. we don't need to close domain.
 	return fn(gs)
 }
+
+func restoreQuery(stmt *ast.BRIEStmt) string {
+	out := bytes.NewBuffer(nil)
+	rc := format.NewRestoreCtx(format.RestoreNameBackQuotes|format.RestoreStringSingleQuotes, out)
+	if err := stmt.Restore(rc); err != nil {
+		return "N/A"
+	}
+	return out.String()
+}

planner/core/planbuilder.go

@@ -3116,16 +3116,22 @@ func buildShowBackupMetaSchema() (*expression.Schema, types.NameSlice) {
 	schema := newColumnsWithNames(len(names))
 	for i := range names {
 		fLen, _ := mysql.GetDefaultFieldLengthAndDecimal(ftypes[i])
+		if ftypes[i] == mysql.TypeVarchar {
+			// the default varchar value is `5`, which might be too short for us.
+			fLen = 255
+		}
 		schema.Append(buildColumnWithName("", names[i], ftypes[i], fLen))
 	}
 	return schema.col2Schema(), schema.names
 }
 
-func buildBRIESchema(kind ast.BRIEKind) (*expression.Schema, types.NameSlice) {
-	if kind == ast.BRIEKindShowBackupMeta {
-		return buildShowBackupMetaSchema()
-	}
+func buildShowBackupQuerySchema() (*expression.Schema, types.NameSlice) {
+	schema := newColumnsWithNames(1)
+	schema.Append(buildColumnWithName("", "Query", mysql.TypeVarchar, 4096))
+	return schema.col2Schema(), schema.names
+}
 
+func buildBackupRestoreSchema(kind ast.BRIEKind) (*expression.Schema, types.NameSlice) {
 	longlongSize, _ := mysql.GetDefaultFieldLengthAndDecimal(mysql.TypeLonglong)
 	datetimeSize, _ := mysql.GetDefaultFieldLengthAndDecimal(mysql.TypeDatetime)
 
@@ -3141,6 +3147,20 @@ func buildBRIESchema(kind ast.BRIEKind) (*expression.Schema, types.NameSlice) {
 	return schema.col2Schema(), schema.names
 }
 
+func buildBRIESchema(kind ast.BRIEKind) (*expression.Schema, types.NameSlice) {
+	switch kind {
+	case ast.BRIEKindShowBackupMeta:
+		return buildShowBackupMetaSchema()
+	case ast.BRIEKindShowQuery:
+		return buildShowBackupQuerySchema()
+	case ast.BRIEKindBackup, ast.BRIEKindRestore:
+		return buildBackupRestoreSchema(kind)
+	default:
+		s := newColumnsWithNames(0)
+		return s.col2Schema(), s.names
+	}
+}
+
 func buildCalibrateResourceSchema() (*expression.Schema, types.NameSlice) {
 	longlongSize, _ := mysql.GetDefaultFieldLengthAndDecimal(mysql.TypeLonglong)
 	schema := newColumnsWithNames(1)
@@ -5240,8 +5260,8 @@ func buildShowSchema(s *ast.ShowStmt, isView bool, isSequence bool) (schema *exp
 		names = []string{"Supported_builtin_functions"}
 		ftypes = []byte{mysql.TypeVarchar}
 	case ast.ShowBackups, ast.ShowRestores:
-		names = []string{"Destination", "State", "Progress", "Queue_time", "Execution_time", "Finish_time", "Connection", "Message"}
-		ftypes = []byte{mysql.TypeVarchar, mysql.TypeVarchar, mysql.TypeDouble, mysql.TypeDatetime, mysql.TypeDatetime, mysql.TypeDatetime, mysql.TypeLonglong, mysql.TypeVarchar}
+		names = []string{"Id", "Destination", "State", "Progress", "Queue_time", "Execution_time", "Finish_time", "Connection", "Message"}
+		ftypes = []byte{mysql.TypeLonglong, mysql.TypeVarchar, mysql.TypeVarchar, mysql.TypeDouble, mysql.TypeDatetime, mysql.TypeDatetime, mysql.TypeDatetime, mysql.TypeLonglong, mysql.TypeVarchar}
 	case ast.ShowPlacementLabels:
 		names = []string{"Key", "Values"}
 		ftypes = []byte{mysql.TypeVarchar, mysql.TypeJSON}