restore: support keyspace rename via --keyspace-mapping flag

pkg/command/restore/cmd.go

@@ -39,6 +39,7 @@ type command struct {
 	dryRun          bool
 	showTables      bool
 	dcMapping       map[string]string
+	ksMapping       map[string]string
 	method          string
 }
 
@@ -94,6 +95,7 @@ func (cmd *command) init() {
 	w.Unwrap().BoolVar(&cmd.dryRun, "dry-run", false, "")
 	w.Unwrap().BoolVar(&cmd.showTables, "show-tables", false, "")
 	w.Unwrap().StringToStringVar(&cmd.dcMapping, "dc-mapping", nil, "")
+	w.Unwrap().StringToStringVar(&cmd.ksMapping, "keyspace-mapping", nil, "")
 	w.Unwrap().StringVar(&cmd.method, "method", "rclone", "")
 }
 
@@ -194,6 +196,13 @@ func (cmd *command) run(args []string) error {
 		props["dc_mapping"] = cmd.dcMapping
 		ok = true
 	}
+	if cmd.Flag("keyspace-mapping").Changed {
+		if cmd.Update() {
+			return wrapper("keyspace-mapping")
+		}
+		props["keyspace_mapping"] = cmd.ksMapping
+		ok = true
+	}
 	if cmd.Flag("method").Changed {
 		props["method"] = cmd.method
 		ok = true

pkg/command/restore/res.yaml

@@ -86,6 +86,19 @@ dc-mapping: |
   Only works with tables restoration (--restore-tables=true). 
   Note: Only DCs that are provided in mappings will be restored.
 
+keyspace-mapping: |
+  Specifies mapping between keyspaces from the backup and keyspaces in the restored(target) cluster.
+  This allows restoring data from a source keyspace into a different (renamed) keyspace.
+
+  The syntax is "source_ks1=target_ks1,source_ks2=target_ks2" where multiple mappings are separated by comma (,)
+  and source and target keyspaces are separated by equal (=).
+
+  Example: "ks1=ks_new" - data from keyspace ks1 in the backup will be restored to keyspace ks_new in the cluster.
+
+  Only works with tables restoration (--restore-tables=true).
+  The target keyspace must already exist in the cluster before starting the restore.
+  Use the '--keyspace' flag to filter which source keyspaces are restored; it still refers to source keyspace names.
+
 method: |
   Control native restore method (See https://manager.docs.scylladb.com/stable/restore/native-restore):
 

pkg/service/restore/index.go

@@ -104,7 +104,7 @@ func (w *tablesWorker) createRemoteDirWorkloads(ctx context.Context, location Lo
 	var rawWorkload []RemoteDirWorkload
 	err := w.forEachManifest(ctx, location, func(m ManifestInfoWithContent) error {
 		return m.ForEachIndexIterWithError(nil, func(fm FilesMeta) error {
-			if !unitsContainTable(w.run.Units, fm.Keyspace, fm.Table) {
+			if !unitsContainTable(w.run.Units, w.target.TargetKeyspace(fm.Keyspace), fm.Table) {
 				return nil
 			}
 
@@ -123,7 +123,7 @@ func (w *tablesWorker) createRemoteDirWorkloads(ctx context.Context, location Lo
 				size += sst.Size
 			}
 			t := TableName{
-				Keyspace: fm.Keyspace,
+				Keyspace: w.target.TargetKeyspace(fm.Keyspace),
 				Table:    fm.Table,
 			}
 			workload := RemoteDirWorkload{

pkg/service/restore/model.go

@@ -23,24 +23,34 @@ import (
 
 // Target specifies what data should be restored and from which locations.
 type Target struct {
-	Location        []backupspec.Location `json:"location"`
-	Keyspace        []string              `json:"keyspace,omitempty"`
-	SnapshotTag     string                `json:"snapshot_tag"`
-	BatchSize       int                   `json:"batch_size,omitempty"`
-	Parallel        int                   `json:"parallel,omitempty"`
-	Transfers       int                   `json:"transfers"`
-	RateLimit       []backup.DCLimit      `json:"rate_limit,omitempty"`
-	AllowCompaction bool                  `json:"allow_compaction,omitempty"`
-	UnpinAgentCPU   bool                  `json:"unpin_agent_cpu"`
-	RestoreSchema   bool                  `json:"restore_schema,omitempty"`
-	RestoreTables   bool                  `json:"restore_tables,omitempty"`
-	Continue        bool                  `json:"continue"`
-	DCMappings      map[string]string     `json:"dc_mapping"`
-	Method          Method                `json:"method,omitempty"`
+	Location         []backupspec.Location `json:"location"`
+	Keyspace         []string              `json:"keyspace,omitempty"`
+	SnapshotTag      string                `json:"snapshot_tag"`
+	BatchSize        int                   `json:"batch_size,omitempty"`
+	Parallel         int                   `json:"parallel,omitempty"`
+	Transfers        int                   `json:"transfers"`
+	RateLimit        []backup.DCLimit      `json:"rate_limit,omitempty"`
+	AllowCompaction  bool                  `json:"allow_compaction,omitempty"`
+	UnpinAgentCPU    bool                  `json:"unpin_agent_cpu"`
+	RestoreSchema    bool                  `json:"restore_schema,omitempty"`
+	RestoreTables    bool                  `json:"restore_tables,omitempty"`
+	Continue         bool                  `json:"continue"`
+	DCMappings       map[string]string     `json:"dc_mapping"`
+	KeyspaceMappings map[string]string     `json:"keyspace_mapping,omitempty"`
+	Method           Method                `json:"method,omitempty"`
 
 	locationInfo []LocationInfo
 }
 
+// TargetKeyspace returns the target keyspace name for the given source keyspace.
+// If no mapping exists for sourceKs, it returns sourceKs unchanged.
+func (t Target) TargetKeyspace(sourceKs string) string {
+	if ks, ok := t.KeyspaceMappings[sourceKs]; ok {
+		return ks
+	}
+	return sourceKs
+}
+
 // Method describes which API should be used by SM during restore.
 type Method string
 
@@ -142,6 +152,17 @@ func (t Target) validateProperties() error {
 	if t.RestoreSchema && t.Method != defaultMethod {
 		return errors.New("restore schema does not support '--method' flag")
 	}
+	if t.RestoreSchema && len(t.KeyspaceMappings) > 0 {
+		return errors.New("restore schema does not support '--keyspace-mapping' flag")
+	}
+	for sourceKs, targetKs := range t.KeyspaceMappings {
+		if sourceKs == "" {
+			return errors.New("keyspace mapping source keyspace must not be empty")
+		}
+		if targetKs == "" {
+			return errors.Errorf("keyspace mapping target keyspace for source %q must not be empty", sourceKs)
+		}
+	}
 	// Check for duplicates in Location
 	allLocations := strset.New()
 	for _, l := range t.Location {

pkg/service/restore/worker.go

@@ -290,6 +290,17 @@ func (w *worker) initTarget(ctx context.Context, t Target, locationInfo []Locati
 		}
 	}
 
+	if len(t.KeyspaceMappings) > 0 {
+		tables, err := w.client.AllTables(ctx)
+		if err != nil {
+			return errors.Wrap(err, "get all tables for keyspace mapping validation")
+		}
+		targetKeyspaces := slices.Collect(maps.Keys(tables))
+		if err := validateKeyspaceMappings(t.KeyspaceMappings, targetKeyspaces); err != nil {
+			return err
+		}
+	}
+
 	w.logger.Info(ctx, "Initialized target", "target", t)
 	return nil
 }
@@ -320,6 +331,23 @@ func (w *worker) validateDCMappings(dcMappings map[string]string, sourceDC, targ
 	return nil
 }
 
+// validateKeyspaceMappings that every target keyspace from mappings exists in target cluster
+// and that target keyspaces are not duplicated.
+func validateKeyspaceMappings(ksMappings map[string]string, targetKeyspaces []string) error {
+	targetKSSet := strset.New(targetKeyspaces...)
+	targetKSMappingSet := strset.New()
+	for _, targetKS := range ksMappings {
+		if !targetKSSet.Has(targetKS) {
+			return errors.Errorf("no such keyspace in target cluster: %s", targetKS)
+		}
+		if targetKSMappingSet.Has(targetKS) {
+			return errors.Errorf("keyspace mapping contains duplicates in target keyspaces: %s", targetKS)
+		}
+		targetKSMappingSet.Add(targetKS)
+	}
+	return nil
+}
+
 func skipRestorePatterns(ctx context.Context, client *scyllaclient.Client, session gocqlx.Session) ([]string, error) {
 	keyspaces, err := client.KeyspacesByType(ctx)
 	if err != nil {
@@ -505,14 +533,15 @@ func (w *worker) initUnits(ctx context.Context, locationInfo []LocationInfo) err
 			foundManifest = true
 
 			filesHandler := func(fm backupspec.FilesMeta) {
-				ru := unitMap[fm.Keyspace]
-				ru.Keyspace = fm.Keyspace
+				targetKs := w.target.TargetKeyspace(fm.Keyspace)
+				ru := unitMap[targetKs]
+				ru.Keyspace = targetKs
 				ru.Size += fm.Size
 
 				for i, t := range ru.Tables {
 					if t.Table == fm.Table {
 						ru.Tables[i].Size += fm.Size
-						unitMap[fm.Keyspace] = ru
+						unitMap[targetKs] = ru
 
 						return
 					}
@@ -522,7 +551,7 @@ func (w *worker) initUnits(ctx context.Context, locationInfo []LocationInfo) err
 					Table: fm.Table,
 					Size:  fm.Size,
 				})
-				unitMap[fm.Keyspace] = ru
+				unitMap[targetKs] = ru
 			}
 
 			return miwc.ForEachIndexIter(w.target.Keyspace, filesHandler)