Implement Raw Restore

cmd/backup.go

@@ -7,6 +7,7 @@ import (
 	"github.com/pingcap/tidb/session"
 	"github.com/spf13/cobra"
 
+	"github.com/pingcap/br/pkg/gluetikv"
 	"github.com/pingcap/br/pkg/summary"
 	"github.com/pingcap/br/pkg/task"
 	"github.com/pingcap/br/pkg/utils"
@@ -21,11 +22,13 @@ func runBackupCommand(command *cobra.Command, cmdName string) error {
 }
 
 func runBackupRawCommand(command *cobra.Command, cmdName string) error {
-	cfg := task.BackupRawConfig{Config: task.Config{LogProgress: HasLogFile()}}
+	cfg := task.BackupRawConfig{
+		RawKvConfig: task.RawKvConfig{Config: task.Config{LogProgress: HasLogFile()}},
+	}
 	if err := cfg.ParseFromFlags(command.Flags()); err != nil {
 		return err
 	}
-	return task.RunBackupRaw(GetDefaultContext(), tidbGlue, cmdName, &cfg)
+	return task.RunBackupRaw(GetDefaultContext(), gluetikv.Glue{}, cmdName, &cfg)
 }
 
 // NewBackupCommand return a full backup subcommand.

cmd/restore.go

@@ -6,6 +6,7 @@ import (
 	"github.com/pingcap/tidb/session"
 	"github.com/spf13/cobra"
 
+	"github.com/pingcap/br/pkg/gluetikv"
 	"github.com/pingcap/br/pkg/summary"
 	"github.com/pingcap/br/pkg/task"
 	"github.com/pingcap/br/pkg/utils"
@@ -19,6 +20,16 @@ func runRestoreCommand(command *cobra.Command, cmdName string) error {
 	return task.RunRestore(GetDefaultContext(), tidbGlue, cmdName, &cfg)
 }
 
+func runRestoreRawCommand(command *cobra.Command, cmdName string) error {
+	cfg := task.RestoreRawConfig{
+		RawKvConfig: task.RawKvConfig{Config: task.Config{LogProgress: HasLogFile()}},
+	}
+	if err := cfg.ParseFromFlags(command.Flags()); err != nil {
+		return err
+	}
+	return task.RunRestoreRaw(GetDefaultContext(), gluetikv.Glue{}, cmdName, &cfg)
+}
+
 // NewRestoreCommand returns a restore subcommand
 func NewRestoreCommand() *cobra.Command {
 	command := &cobra.Command{
@@ -43,6 +54,7 @@ func NewRestoreCommand() *cobra.Command {
 		newFullRestoreCommand(),
 		newDbRestoreCommand(),
 		newTableRestoreCommand(),
+		newRawRestoreCommand(),
 	)
 	task.DefineRestoreFlags(command.PersistentFlags())
 
@@ -83,3 +95,16 @@ func newTableRestoreCommand() *cobra.Command {
 	task.DefineTableFlags(command)
 	return command
 }
+
+func newRawRestoreCommand() *cobra.Command {
+	command := &cobra.Command{
+		Use:   "raw",
+		Short: "(experimental) restore a raw kv range to TiKV cluster",
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return runRestoreRawCommand(cmd, "Raw restore")
+		},
+	}
+
+	task.DefineRawRestoreFlags(command)
+	return command
+}

pkg/gluetidb/glue.go

@@ -8,19 +8,20 @@ import (
 
 	"github.com/pingcap/parser/model"
 	pd "github.com/pingcap/pd/client"
-	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/executor"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/session"
-	"github.com/pingcap/tidb/store/tikv"
 
 	"github.com/pingcap/br/pkg/glue"
+	"github.com/pingcap/br/pkg/gluetikv"
 )
 
 // Glue is an implementation of glue.Glue using a new TiDB session.
-type Glue struct{}
+type Glue struct {
+	tikvGlue gluetikv.Glue
+}
 
 type tidbSession struct {
 	se session.Session
@@ -41,15 +42,8 @@ func (Glue) CreateSession(store kv.Storage) (glue.Session, error) {
 }
 
 // Open implements glue.Glue
-func (Glue) Open(path string, option pd.SecurityOption) (kv.Storage, error) {
-	if option.CAPath != "" {
-		conf := config.GetGlobalConfig()
-		conf.Security.ClusterSSLCA = option.CAPath
-		conf.Security.ClusterSSLCert = option.CertPath
-		conf.Security.ClusterSSLKey = option.KeyPath
-		config.StoreGlobalConfig(conf)
-	}
-	return tikv.Driver{}.Open(path)
+func (g Glue) Open(path string, option pd.SecurityOption) (kv.Storage, error) {
+	return g.tikvGlue.Open(path, option)
 }
 
 // OwnsStorage implements glue.Glue

pkg/gluetikv/glue.go

@@ -0,0 +1,41 @@
+package gluetikv
+
+import (
+	pd "github.com/pingcap/pd/client"
+	"github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/domain"
+	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/store/tikv"
+
+	"github.com/pingcap/br/pkg/glue"
+)
+
+// Glue is an implementation of glue.Glue that accesses only TiKV without TiDB.
+type Glue struct{}
+
+// GetDomain implements glue.Glue
+func (Glue) GetDomain(store kv.Storage) (*domain.Domain, error) {
+	return nil, nil
+}
+
+// CreateSession implements glue.Glue
+func (Glue) CreateSession(store kv.Storage) (glue.Session, error) {
+	return nil, nil
+}
+
+// Open implements glue.Glue
+func (Glue) Open(path string, option pd.SecurityOption) (kv.Storage, error) {
+	if option.CAPath != "" {
+		conf := config.GetGlobalConfig()
+		conf.Security.ClusterSSLCA = option.CAPath
+		conf.Security.ClusterSSLCert = option.CertPath
+		conf.Security.ClusterSSLKey = option.KeyPath
+		config.StoreGlobalConfig(conf)
+	}
+	return tikv.Driver{}.Open(path)
+}
+
+// OwnsStorage implements glue.Glue
+func (Glue) OwnsStorage() bool {
+	return true
+}

pkg/restore/client.go

@@ -3,8 +3,10 @@
 package restore
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
+	"encoding/hex"
 	"encoding/json"
 	"math"
 	"sort"
@@ -96,33 +98,96 @@ func (rc *Client) IsOnline() bool {
 
 // Close a client
 func (rc *Client) Close() {
-	rc.db.Close()
+	// rc.db can be nil in raw kv mode.
+	if rc.db != nil {
+		rc.db.Close()
+	}
 	rc.cancel()
 	log.Info("Restore client closed")
 }
 
 // InitBackupMeta loads schemas from BackupMeta to initialize RestoreClient
 func (rc *Client) InitBackupMeta(backupMeta *backup.BackupMeta, backend *backup.StorageBackend) error {
-	databases, err := utils.LoadBackupTables(backupMeta)
-	if err != nil {
-		return errors.Trace(err)
+	if !backupMeta.IsRawKv {
+		databases, err := utils.LoadBackupTables(backupMeta)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		rc.databases = databases
 	}
 	var ddlJobs []*model.Job
-	err = json.Unmarshal(backupMeta.GetDdls(), &ddlJobs)
+	err := json.Unmarshal(backupMeta.GetDdls(), &ddlJobs)
 	if err != nil {
 		return errors.Trace(err)
 	}
-	rc.databases = databases
 	rc.ddlJobs = ddlJobs
 	rc.backupMeta = backupMeta
 	log.Info("load backupmeta", zap.Int("databases", len(rc.databases)), zap.Int("jobs", len(rc.ddlJobs)))
 
 	metaClient := NewSplitClient(rc.pdClient, rc.tlsConf)
 	importClient := NewImportClient(metaClient, rc.tlsConf)
-	rc.fileImporter = NewFileImporter(rc.ctx, metaClient, importClient, backend, rc.rateLimit)
+	rc.fileImporter = NewFileImporter(rc.ctx, metaClient, importClient, backend, backupMeta.IsRawKv, rc.rateLimit)
 	return nil
 }
 
+// IsRawKvMode checks whether the backup data is in raw kv format, in which case transactional recover is forbidden.
+func (rc *Client) IsRawKvMode() bool {
+	return rc.backupMeta.IsRawKv
+}
+
+// GetFilesInRawRange gets all files that are in the given range or intersects with the given range.
+func (rc *Client) GetFilesInRawRange(startKey []byte, endKey []byte, cf string) ([]*backup.File, error) {
+	if !rc.IsRawKvMode() {
+		return nil, errors.New("the backup data is not in raw kv mode")
+	}
+
+	for _, rawRange := range rc.backupMeta.RawRanges {
+		// First check whether the given range is backup-ed. If not, we cannot perform the restore.
+		if rawRange.Cf != cf {
+			continue
+		}
+
+		if (len(rawRange.EndKey) > 0 && bytes.Compare(startKey, rawRange.EndKey) >= 0) ||
+			(len(endKey) > 0 && bytes.Compare(rawRange.StartKey, endKey) >= 0) {
+			// The restoring range is totally out of the current range. Skip it.
+			continue
+		}
+
+		if bytes.Compare(startKey, rawRange.StartKey) < 0 ||
+			utils.CompareEndKey(endKey, rawRange.EndKey) > 0 {
+			// Only partial of the restoring range is in the current backup-ed range. So the given range can't be fully
+			// restored.
+			return nil, errors.New("no backup data in the range")
+		}
+
+		// We have found the range that contains the given range. Find all necessary files.
+		files := make([]*backup.File, 0)
+
+		for _, file := range rc.backupMeta.Files {
+			if file.Cf != cf {
+				continue
+			}
+
+			if len(file.EndKey) > 0 && bytes.Compare(file.EndKey, startKey) < 0 {
+				// The file is before the range to be restored.
+				continue
+			}
+			if len(endKey) > 0 && bytes.Compare(endKey, file.StartKey) <= 0 {
+				// The file is after the range to be restored.
+				// The specified endKey is exclusive, so when it equals to a file's startKey, the file is still skipped.
+				continue
+			}
+
+			files = append(files, file)
+		}
+
+		// There should be at most one backed up range that covers the restoring range.
+		return files, nil
+	}
+
+	return nil, errors.New("no backup data in the range")
+}
+
 // SetConcurrency sets the concurrency of dbs tables files
 func (rc *Client) SetConcurrency(c uint) {
 	rc.workerPool = utils.NewWorkerPool(c, "file")
@@ -323,6 +388,63 @@ func (rc *Client) RestoreFiles(
 	return nil
 }
 
+// RestoreRaw tries to restore raw keys in the specified range.
+func (rc *Client) RestoreRaw(startKey []byte, endKey []byte, files []*backup.File, updateCh chan<- struct{}) error {
+	start := time.Now()
+	defer func() {
+		elapsed := time.Since(start)
+		log.Info("Restore Raw",
+			zap.String("startKey", hex.EncodeToString(startKey)),
+			zap.String("endKey", hex.EncodeToString(endKey)),
+			zap.Duration("take", elapsed))
+	}()
+	errCh := make(chan error, len(rc.databases))
+	wg := new(sync.WaitGroup)
+	defer close(errCh)
+
+	err := rc.fileImporter.SetRawRange(startKey, endKey)
+	if err != nil {
+
+		return errors.Trace(err)
+	}
+
+	emptyRules := &RewriteRules{}
+	for _, file := range files {
+		wg.Add(1)
+		fileReplica := file
+		rc.workerPool.Apply(
+			func() {
+				defer wg.Done()
+				select {
+				case <-rc.ctx.Done():
+					errCh <- nil
+				case errCh <- rc.fileImporter.Import(fileReplica, emptyRules):
+					updateCh <- struct{}{}
+				}
+			})
+	}
+	for range files {
+		err := <-errCh
+		if err != nil {
+			rc.cancel()
+			wg.Wait()
+			log.Error(
+				"restore raw range failed",
+				zap.String("startKey", hex.EncodeToString(startKey)),
+				zap.String("endKey", hex.EncodeToString(endKey)),
+				zap.Error(err),
+			)
+			return err
+		}
+	}
+	log.Info(
+		"finish to restore raw range",
+		zap.String("startKey", hex.EncodeToString(startKey)),
+		zap.String("endKey", hex.EncodeToString(endKey)),
+	)
+	return nil
+}
+
 //SwitchToImportMode switch tikv cluster to import mode
 func (rc *Client) SwitchToImportMode(ctx context.Context) error {
 	return rc.switchTiKVMode(ctx, import_sstpb.SwitchMode_Import)

pkg/restore/db.go

@@ -29,6 +29,10 @@ func NewDB(g glue.Glue, store kv.Storage) (*DB, error) {
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
+	// The session may be nil in raw kv mode
+	if se == nil {
+		return nil, nil
+	}
 	// Set SQL mode to None for avoiding SQL compatibility problem
 	err = se.Execute(context.Background(), "set @@sql_mode=''")
 	if err != nil {