From db3309d175a9b4914767901671d78ee732e866cf Mon Sep 17 00:00:00 2001 From: Vasil Averyanau Date: Thu, 16 Jan 2025 16:33:29 +0100 Subject: [PATCH] feat(restore): adds --dc-mapping flag to restore command This adds support for --dc-mapping flag to restore command. It specifies mapping between DCs from the backup and DCs in the restored(target) cluster. All DCs from the source cluster should be explicitly mapped to all DCs in the target cluster. The only exception is when source and target cluster has exact match: source dcs == target dcs. Only works with tables restoration (--restore-tables=true). Syntax: "source_dc1,source_dc2=>target_dc1,target_dc2" Multiple mappings are separated by semicolons (;). Exclamation mark (!) before a DC indicates that it should be ignored during restore. Examples: "dc1,dc2=>dc3" - data from dc1 and dc2 DCs should be restored to dc3 DC. "dc1,dc2=>dc3,!dc4" - data from dc1 and dc2 DCs should be restored to dc3 DC. Ignoring dc4 DC from target cluster. "dc1,!dc2=>dc2" - data from dc1 should be restored to dc2 DC. Ignoring dc2 from source cluster. Fixes: #3829 --- pkg/command/restore/cmd.go | 9 ++ pkg/command/restore/dcmappings.go | 87 ++++++++++++ pkg/command/restore/dcmappings_test.go | 185 +++++++++++++++++++++++++ pkg/command/restore/res.yaml | 13 ++ 4 files changed, 294 insertions(+) create mode 100644 pkg/command/restore/dcmappings.go create mode 100644 pkg/command/restore/dcmappings_test.go diff --git a/pkg/command/restore/cmd.go b/pkg/command/restore/cmd.go index 834f132a8..58db403ce 100644 --- a/pkg/command/restore/cmd.go +++ b/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) diff --git a/pkg/command/restore/dcmappings.go b/pkg/command/restore/dcmappings.go new file mode 100644 index 000000000..43932a9b3 --- /dev/null +++ b/pkg/command/restore/dcmappings.go @@ -0,0 +1,87 @@ +package restore + +import ( + "slices" + "strings" + + "github.com/pkg/errors" +) + +type dcMappings []dcMapping + +type dcMapping struct { + Source []string `json:"source"` + IgnoreSource []string `json:"ignore_source"` + Target []string `json:"target"` + IgnoreTarget []string `json:"ignore_target"` +} + +const ignoreDCPrefix = "!" + +// Set parses --dc-mapping flag, where the syntax is following: +// ; - used to split different mappings +// => - used to split source => target DCs +// , - used to seprate DCs +// ! - used to ignore DC. +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, mapping.IgnoreSource = parseDCList(strings.Split(sourceTargetParts[0], ",")) + mapping.Target, mapping.IgnoreTarget = parseDCList(strings.Split(sourceTargetParts[1], ",")) + + *dcm = append(*dcm, mapping) + } + return nil +} + +func parseDCList(list []string) (dcs, ignore []string) { + for _, dc := range list { + if strings.HasPrefix(dc, ignoreDCPrefix) { + ignore = append(ignore, strings.TrimPrefix(dc, ignoreDCPrefix)) + continue + } + dcs = append(dcs, dc) + } + return dcs, ignore +} + +// 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 { + source := slices.Concat(mapping.Source, addIgnorePrefix(mapping.IgnoreSource)) + target := slices.Concat(mapping.Target, addIgnorePrefix(mapping.IgnoreTarget)) + res.WriteString( + strings.Join(source, ",") + "=>" + strings.Join(target, ","), + ) + if i != len(*dcm)-1 { + res.WriteString(";") + } + } + return res.String() +} + +func addIgnorePrefix(ignore []string) []string { + var result []string + for _, v := range ignore { + result = append(result, ignoreDCPrefix+v) + } + return result +} + +// Type implements pflag.Value interface +func (dcm *dcMappings) Type() string { + return "dc-mapping" +} diff --git a/pkg/command/restore/dcmappings_test.go b/pkg/command/restore/dcmappings_test.go new file mode 100644 index 000000000..8697d5aa6 --- /dev/null +++ b/pkg/command/restore/dcmappings_test.go @@ -0,0 +1,185 @@ +// 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: []string{"dc1"}, Target: []string{"dc2"}}, + }, + }, + { + input: "dc1, dc2=>dc1, dc2", + expectedMappings: dcMappings{ + {Source: []string{"dc1", "dc2"}, Target: []string{"dc1", "dc2"}}, + }, + }, + { + input: "dc1=>dc3;dc2=>dc4", + expectedMappings: dcMappings{ + {Source: []string{"dc1"}, Target: []string{"dc3"}}, + {Source: []string{"dc2"}, Target: []string{"dc4"}}, + }, + }, + { + input: "dc1,dc2=>dc3", + expectedMappings: dcMappings{ + {Source: []string{"dc1", "dc2"}, Target: []string{"dc3"}}, + }, + }, + { + input: "dc1,!dc2=>dc3", + expectedMappings: dcMappings{ + { + Source: []string{"dc1"}, + Target: []string{"dc3"}, + IgnoreSource: []string{"dc2"}, + }, + }, + }, + { + input: "dc1,!dc2=>dc3,!dc4", + expectedMappings: dcMappings{ + { + Source: []string{"dc1"}, + Target: []string{"dc3"}, + IgnoreSource: []string{"dc2"}, + IgnoreTarget: []string{"dc4"}, + }, + }, + }, + { + input: "dc1,!dc2=>dc3,!dc4", + expectedMappings: dcMappings{ + { + Source: []string{"dc1"}, + Target: []string{"dc3"}, + + IgnoreSource: []string{"dc2"}, + IgnoreTarget: []string{"dc4"}, + }, + }, + }, + { + input: "!dc1,dc2=>dc3,!dc4", + expectedMappings: dcMappings{ + { + Source: []string{"dc2"}, + Target: []string{"dc3"}, + + IgnoreSource: []string{"dc1"}, + IgnoreTarget: []string{"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) + } + slices.EqualFunc(tc.expectedMappings, mappings, func(a, b dcMapping) bool { + return slices.Equal(a.Source, b.Source) && + slices.Equal(a.Target, b.Target) && + slices.Equal(a.IgnoreSource, b.IgnoreSource) && + slices.Equal(a.IgnoreTarget, b.IgnoreTarget) + }) + }) + } + +} + +func TestDCMappingString(t *testing.T) { + testCases := []struct { + mappings dcMappings + expected string + }{ + { + mappings: dcMappings{ + {Source: []string{"dc1"}, Target: []string{"dc2"}}, + }, + expected: "dc1=>dc2", + }, + { + mappings: dcMappings{ + {Source: []string{"dc1"}, Target: []string{"dc2"}}, + {Source: []string{"dc3"}, Target: []string{"dc4"}}, + }, + expected: "dc1=>dc2;dc3=>dc4", + }, + { + mappings: dcMappings{ + { + Source: []string{"dc1"}, + Target: []string{"dc2"}, + IgnoreSource: []string{"dc2"}, + }, + }, + expected: "dc1,!dc2=>dc2", + }, + { + mappings: dcMappings{ + { + Source: []string{"dc1"}, + Target: []string{"dc2"}, + IgnoreSource: []string{"dc2"}, + IgnoreTarget: []string{"dc3"}, + }, + }, + expected: "dc1,!dc2=>dc2,!dc3", + }, + {}, + } + + 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) + } + }) + } +} diff --git a/pkg/command/restore/res.yaml b/pkg/command/restore/res.yaml index eaeb4bd56..bfe58fa30 100644 --- a/pkg/command/restore/res.yaml +++ b/pkg/command/restore/res.yaml @@ -72,3 +72,16 @@ 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. + All DCs from source cluster should be explicitly mapped to all DCs in the target cluster. The only exception is when + source and target cluster has exact match: source dcs == target dcs. + Only works with tables restoration (--restore-tables=true). + Syntax: + "source_dc1,source_dc2=>target_dc1,target_dc2" + Multiple mappings are separated by semicolons (;). Exclamation mark (!) before a DC indicates that it should be ignored during restore. + Examples: + "dc1,dc2=>dc3" - data from dc1 and dc2 DCs should be restored to dc3 DC. + "dc1,dc2=>dc3,!dc4" - data from dc1 and dc2 DCs should be restored to dc3 DC. Ignoring dc4 DC from target cluster. + "dc1,!dc2=>dc2" - data from dc1 should be restored to dc2 DC. Ignoring dc2 from source cluster.