feat(restore): adds --dc-mapping flag to restore command

docs/source/sctool/partials/sctool_restore.yaml

@@ -26,6 +26,8 @@ options:
       usage: |
         Task schedule as a cron `expression`.
         It supports the extended syntax including @monthly, @weekly, @daily, @midnight, @hourly, @every X[h|m|s].
+    - name: dc-mapping
+      usage: "Specifies mapping between DCs from the backup and DCs in the restored(target) cluster.\n\nThe Syntax is \"source_dc1=>target_dc1;source_dc2=>target_dc2\" where multiple mappings are separated by semicolons (;)\nand source and target DCs are separated by arrow (=>).\n\nExample: \"dc1=>dc3;dc2=>dc4\" - data from dc1 should be restored to dc3 and data from dc2 should be restored to dc4.\n\nOnly works with tables restoration (--restore-tables=true). \nNote: Only DCs that are provided in mappings will be restored.\n"
     - name: dry-run
       default_value: "false"
       usage: |
@@ -90,6 +92,7 @@ options:
         The `<dc>` parameter is optional. It allows you to specify the datacenter whose nodes will be used to restore the data
         from this location in a multi-dc setting, it must match Scylla nodes datacenter.
         By default, all live nodes are used to restore data from specified locations.
+        If `--dc-mapping` is used, then `<dc>` parameter will be ignored.
 
         Note that specifying datacenters closest to backup locations might reduce download time of restored data.
         The supported storage '<provider>'s are 'azure', 'gcs', 's3'.

docs/source/sctool/partials/sctool_restore_update.yaml

@@ -24,6 +24,8 @@ options:
       usage: |
         Task schedule as a cron `expression`.
         It supports the extended syntax including @monthly, @weekly, @daily, @midnight, @hourly, @every X[h|m|s].
+    - name: dc-mapping
+      usage: "Specifies mapping between DCs from the backup and DCs in the restored(target) cluster.\n\nThe Syntax is \"source_dc1=>target_dc1;source_dc2=>target_dc2\" where multiple mappings are separated by semicolons (;)\nand source and target DCs are separated by arrow (=>).\n\nExample: \"dc1=>dc3;dc2=>dc4\" - data from dc1 should be restored to dc3 and data from dc2 should be restored to dc4.\n\nOnly works with tables restoration (--restore-tables=true). \nNote: Only DCs that are provided in mappings will be restored.\n"
     - name: dry-run
       default_value: "false"
       usage: |
@@ -88,6 +90,7 @@ options:
         The `<dc>` parameter is optional. It allows you to specify the datacenter whose nodes will be used to restore the data
         from this location in a multi-dc setting, it must match Scylla nodes datacenter.
         By default, all live nodes are used to restore data from specified locations.
+        If `--dc-mapping` is used, then `<dc>` parameter will be ignored.
 
         Note that specifying datacenters closest to backup locations might reduce download time of restored data.
         The supported storage '<provider>'s are 'azure', 'gcs', 's3'.

pkg/command/restore/cmd.go

@@ -37,6 +37,7 @@ type command struct {
 	restoreTables   bool
 	dryRun          bool
 	showTables      bool
+	dcMapping       dcMappings
 }
 
 func NewCommand(client *managerclient.Client) *cobra.Command {
@@ -90,6 +91,7 @@ func (cmd *command) init() {
 	w.Unwrap().BoolVar(&cmd.restoreTables, "restore-tables", false, "")
 	w.Unwrap().BoolVar(&cmd.dryRun, "dry-run", false, "")
 	w.Unwrap().BoolVar(&cmd.showTables, "show-tables", false, "")
+	w.Unwrap().Var(&cmd.dcMapping, "dc-mapping", "")
 }
 
 func (cmd *command) run(args []string) error {
@@ -182,6 +184,13 @@ func (cmd *command) run(args []string) error {
 		props["restore_tables"] = cmd.restoreTables
 		ok = true
 	}
+	if cmd.Flag("dc-mapping").Changed {
+		if cmd.Update() {
+			return wrapper("dc-mapping")
+		}
+		props["dc_mapping"] = cmd.dcMapping
+		ok = true
+	}
 
 	if cmd.dryRun {
 		res, err := cmd.client.GetRestoreTarget(cmd.Context(), cmd.cluster, task)

pkg/command/restore/dcmappings.go

@@ -0,0 +1,59 @@
+// Copyright (C) 2025 ScyllaDB
+
+package restore
+
+import (
+	"strings"
+
+	"github.com/pkg/errors"
+)
+
+type dcMappings []dcMapping
+
+type dcMapping struct {
+	Source string `json:"source"`
+	Target string `json:"target"`
+}
+
+// Set parses --dc-mapping flag, where the syntax is following:
+// ; - used to split different mappings
+// => - used to split source => target DCs.
+func (dcm *dcMappings) Set(v string) error {
+	mappingParts := strings.Split(v, ";")
+	for _, dcMapPart := range mappingParts {
+		sourceTargetParts := strings.Split(dcMapPart, "=>")
+		if len(sourceTargetParts) != 2 {
+			return errors.New("invalid syntax, mapping should be in a format of sourceDcs=>targetDcs, but got: " + dcMapPart)
+		}
+		if sourceTargetParts[0] == "" || sourceTargetParts[1] == "" {
+			return errors.New("invalid syntax, mapping should be in a format of sourceDcs=>targetDcs, but got: " + dcMapPart)
+		}
+
+		var mapping dcMapping
+		mapping.Source = strings.TrimSpace(sourceTargetParts[0])
+		mapping.Target = strings.TrimSpace(sourceTargetParts[1])
+
+		*dcm = append(*dcm, mapping)
+	}
+	return nil
+}
+
+// String builds --dc-mapping flag back from struct.
+func (dcm *dcMappings) String() string {
+	if dcm == nil {
+		return ""
+	}
+	var res strings.Builder
+	for i, mapping := range *dcm {
+		res.WriteString(mapping.Source + "=>" + mapping.Target)
+		if i != len(*dcm)-1 {
+			res.WriteString(";")
+		}
+	}
+	return res.String()
+}
+
+// Type implements pflag.Value interface.
+func (dcm *dcMappings) Type() string {
+	return "dc-mapping"
+}

pkg/command/restore/dcmappings_test.go

@@ -0,0 +1,110 @@
+// Copyright (C) 2025 ScyllaDB
+package restore
+
+import (
+	"fmt"
+	"slices"
+	"testing"
+)
+
+func TestSetDCMapping(t *testing.T) {
+	testCases := []struct {
+		input            string
+		expectedErr      bool
+		expectedMappings dcMappings
+	}{
+		{
+			input: "dc1=>dc2",
+			expectedMappings: dcMappings{
+				{Source: "dc1", Target: "dc2"},
+			},
+		},
+		{
+			input: " dc1 => dc1 ",
+			expectedMappings: dcMappings{
+				{Source: "dc1", Target: "dc1"},
+			},
+		},
+		{
+			input: "dc1=>dc3;dc2=>dc4",
+			expectedMappings: dcMappings{
+				{Source: "dc1", Target: "dc3"},
+				{Source: "dc2", Target: "dc4"},
+			},
+		},
+		{
+			input:            "dc1=>dc3=>dc2=>dc4",
+			expectedMappings: dcMappings{},
+			expectedErr:      true,
+		},
+		{
+			input:            ";",
+			expectedMappings: dcMappings{},
+			expectedErr:      true,
+		},
+		{
+			input:            "=>",
+			expectedMappings: dcMappings{},
+			expectedErr:      true,
+		},
+		{
+			input:            "dc1=>",
+			expectedMappings: dcMappings{},
+			expectedErr:      true,
+		},
+		{
+			input:            "dc1=>;",
+			expectedMappings: dcMappings{},
+			expectedErr:      true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.input, func(t *testing.T) {
+			var mappings dcMappings
+
+			err := mappings.Set(tc.input)
+			if tc.expectedErr && err == nil {
+				t.Fatal("Expected err, but got nil")
+			}
+			if !tc.expectedErr && err != nil {
+				t.Fatalf("Unexpected err: %v", err)
+			}
+			if !slices.Equal(mappings, tc.expectedMappings) {
+				t.Fatalf("Expected %v, but got %v", tc.expectedMappings, mappings)
+			}
+		})
+	}
+
+}
+
+func TestDCMappingString(t *testing.T) {
+	testCases := []struct {
+		mappings dcMappings
+		expected string
+	}{
+		{
+			mappings: dcMappings{
+				{Source: "dc1", Target: "dc2"},
+			},
+			expected: "dc1=>dc2",
+		},
+		{
+			mappings: dcMappings{
+				{Source: "dc1", Target: "dc2"},
+				{Source: "dc3", Target: "dc4"},
+			},
+			expected: "dc1=>dc2;dc3=>dc4",
+		},
+		{},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			actual := tc.mappings.String()
+			if actual != tc.expected {
+				t.Fatalf("Expected %q, but got %q", tc.expected, actual)
+			}
+		})
+	}
+}

pkg/command/restore/res.yaml

@@ -15,6 +15,7 @@ location: |
   The `<dc>` parameter is optional. It allows you to specify the datacenter whose nodes will be used to restore the data
   from this location in a multi-dc setting, it must match Scylla nodes datacenter.
   By default, all live nodes are used to restore data from specified locations.
+  If `--dc-mapping` is used, then `<dc>` parameter will be ignored.
 
   Note that specifying datacenters closest to backup locations might reduce download time of restored data.
   The supported storage '<provider>'s are 'azure', 'gcs', 's3'.
@@ -72,3 +73,14 @@ dry-run: |
 
 show-tables: |
   Prints table names together with keyspace, used in combination with --dry-run.
+
+dc-mapping: |
+  Specifies mapping between DCs from the backup and DCs in the restored(target) cluster.
+
+  The Syntax is "source_dc1=>target_dc1;source_dc2=>target_dc2" where multiple mappings are separated by semicolons (;)
+  and source and target DCs are separated by arrow (=>).
+
+  Example: "dc1=>dc3;dc2=>dc4" - data from dc1 should be restored to dc3 and data from dc2 should be restored to dc4.
+
+  Only works with tables restoration (--restore-tables=true). 
+  Note: Only DCs that are provided in mappings will be restored.

pkg/service/restore/batch.go

@@ -57,7 +57,7 @@ type batchDispatcher struct {
 	hostShardCnt map[string]uint
 }
 
-func newBatchDispatcher(workload Workload, batchSize int, hostShardCnt map[string]uint, locationHosts map[backupspec.Location][]string) *batchDispatcher {
+func newBatchDispatcher(workload Workload, batchSize int, hostShardCnt map[string]uint, locationInfo []LocationInfo) *batchDispatcher {
 	sortWorkload(workload)
 	var shards uint
 	for _, sh := range hostShardCnt {
@@ -70,7 +70,7 @@ func newBatchDispatcher(workload Workload, batchSize int, hostShardCnt map[strin
 		mu:                    sync.Mutex{},
 		wait:                  make(chan struct{}),
 		workload:              workload,
-		workloadProgress:      newWorkloadProgress(workload, locationHosts),
+		workloadProgress:      newWorkloadProgress(workload, locationInfo),
 		batchSize:             batchSize,
 		expectedShardWorkload: workload.TotalSize / int64(shards),
 		hostShardCnt:          hostShardCnt,
@@ -106,22 +106,22 @@ type remoteSSTableDirProgress struct {
 	RemainingSSTables []RemoteSSTable
 }
 
-func newWorkloadProgress(workload Workload, locationHosts map[backupspec.Location][]string) workloadProgress {
+func newWorkloadProgress(workload Workload, locationInfo []LocationInfo) workloadProgress {
 	dcBytes := make(map[string]int64)
-	locationDC := make(map[string][]string)
 	p := make([]remoteSSTableDirProgress, len(workload.RemoteDir))
 	for i, rdw := range workload.RemoteDir {
 		dcBytes[rdw.DC] += rdw.Size
-		locationDC[rdw.Location.StringWithoutDC()] = append(locationDC[rdw.Location.StringWithoutDC()], rdw.DC)
 		p[i] = remoteSSTableDirProgress{
 			RemainingSize:     rdw.Size,
 			RemainingSSTables: rdw.SSTables,
 		}
 	}
-	hostDCAccess := make(map[string][]string)
-	for loc, hosts := range locationHosts {
-		for _, h := range hosts {
-			hostDCAccess[h] = append(hostDCAccess[h], locationDC[loc.StringWithoutDC()]...)
+	hostDCAccess := map[string][]string{}
+	for _, l := range locationInfo {
+		for dc, hosts := range l.DCHosts {
+			for _, h := range hosts {
+				hostDCAccess[h] = append(hostDCAccess[h], dc)
+			}
 		}
 	}
 	return workloadProgress{
@@ -201,8 +201,8 @@ func (bd *batchDispatcher) ValidateAllDispatched() error {
 	for i, rdp := range bd.workloadProgress.remoteDir {
 		if rdp.RemainingSize != 0 || len(rdp.RemainingSSTables) != 0 {
 			rdw := bd.workload.RemoteDir[i]
-			return errors.Errorf("failed to restore sstables from location %s table %s.%s (%d bytes). See logs for more info",
-				rdw.Location, rdw.Keyspace, rdw.Table, rdw.Size)
+			return errors.Errorf("failed to restore sstables from location %s dc %s table %s.%s (%d bytes). See logs for more info",
+				rdw.Location, rdw.DC, rdw.Keyspace, rdw.Table, rdw.Size)
 		}
 	}
 	for dc, bytes := range bd.workloadProgress.dcBytesToBeRestored {
@@ -257,7 +257,7 @@ func (bd *batchDispatcher) dispatchBatch(host string) (batch, bool) {
 		if slices.Contains(bd.workloadProgress.hostFailedDC[host], rdw.DC) {
 			continue
 		}
-		// Sip dir from location without access
+		// Skip dir from location without access
 		if !slices.Contains(bd.workloadProgress.hostDCAccess[host], rdw.DC) {
 			continue
 		}