diff --git a/pkg/ccl/cliccl/debug_backup.go b/pkg/ccl/cliccl/debug_backup.go index 6941550a9f22..1e1f4fab6a3e 100644 --- a/pkg/ccl/cliccl/debug_backup.go +++ b/pkg/ccl/cliccl/debug_backup.go @@ -61,6 +61,25 @@ var debugBackupArgs struct { destination string format string nullas string + maxRows int + startKey cli.MVCCKey + + rowCount int +} + +// setDebugBackupArgsDefault set the default values in debugBackupArgs. +// This function is called in every test that exercises debug backup +// command-line parsing. +func setDebugContextDefault() { + debugBackupArgs.externalIODir = "" + debugBackupArgs.exportTableName = "" + debugBackupArgs.readTime = "" + debugBackupArgs.destination = "" + debugBackupArgs.format = "csv" + debugBackupArgs.nullas = "null" + debugBackupArgs.maxRows = 0 + debugBackupArgs.startKey = cli.MVCCKey{} + debugBackupArgs.rowCount = 0 } func init() { @@ -150,6 +169,17 @@ func init() { "null", /*value*/ cliflags.ExportCSVNullas.Usage()) + exportDataCmd.Flags().IntVar( + &debugBackupArgs.maxRows, + cliflags.MaxRows.Name, + 0, + cliflags.MaxRows.Usage()) + + exportDataCmd.Flags().Var( + &debugBackupArgs.startKey, + cliflags.StartKey.Name, + cliflags.StartKey.Usage()) + cli.DebugCmd.AddCommand(backupCmds) backupSubCmds := []*cobra.Command{ @@ -305,7 +335,6 @@ func runListIncrementalCmd(cmd *cobra.Command, args []string) error { } func runExportDataCmd(cmd *cobra.Command, args []string) error { - if debugBackupArgs.exportTableName == "" { return errors.New("export data requires table name specified by --table flag") } @@ -391,6 +420,9 @@ func showData( if err := processEntryFiles(ctx, rf, files, entry.Span, endTime, writer); err != nil { return err } + if debugBackupArgs.maxRows != 0 && debugBackupArgs.rowCount >= debugBackupArgs.maxRows { + break + } } if debugBackupArgs.destination != "" { @@ -501,6 +533,9 @@ func processEntryFiles( defer iter.Close() startKeyMVCC, endKeyMVCC := storage.MVCCKey{Key: span.Key}, storage.MVCCKey{Key: span.EndKey} + if len(debugBackupArgs.startKey.Key) != 0 { + startKeyMVCC.Key = debugBackupArgs.startKey.Key + } kvFetcher := row.MakeBackupSSTKVFetcher(startKeyMVCC, endKeyMVCC, iter, endTime) if err := rf.StartScanFrom(ctx, &kvFetcher); err != nil { @@ -527,6 +562,13 @@ func processEntryFiles( return err } writer.Flush() + + if debugBackupArgs.maxRows != 0 { + debugBackupArgs.rowCount++ + if debugBackupArgs.rowCount >= debugBackupArgs.maxRows { + break + } + } } return nil } diff --git a/pkg/ccl/cliccl/debug_backup_test.go b/pkg/ccl/cliccl/debug_backup_test.go index b4e15d030a27..6f8b6acd9233 100644 --- a/pkg/ccl/cliccl/debug_backup_test.go +++ b/pkg/ccl/cliccl/debug_backup_test.go @@ -62,6 +62,7 @@ func TestShowSummary(t *testing.T) { sqlDB.Exec(t, `BACKUP DATABASE testDB TO $1 AS OF SYSTEM TIME `+ts2.AsOfSystemTime(), backupPath) t.Run("show-summary-without-types-or-tables", func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup show %s --external-io-dir=%s", dbOnlyBackupPath, dir)) require.NoError(t, err) expectedOutput := fmt.Sprintf( @@ -91,6 +92,7 @@ func TestShowSummary(t *testing.T) { }) t.Run("show-summary-with-full-information", func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup show %s --external-io-dir=%s", backupPath, dir)) require.NoError(t, err) @@ -172,6 +174,7 @@ func TestListBackups(t *testing.T) { sqlDB.Exec(t, fmt.Sprintf(`BACKUP DATABASE testDB INTO $1 AS OF SYSTEM TIME '%s'`, ts[2].AsOfSystemTime()), backupPath) t.Run("show-backups-with-backups-in-collection", func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup list-backups %s --external-io-dir=%s", backupPath, dir)) require.NoError(t, err) @@ -264,6 +267,10 @@ func TestExportData(t *testing.T) { ts1 := hlc.Timestamp{WallTime: timeutil.Now().UnixNano()} sqlDB.Exec(t, fmt.Sprintf(`BACKUP TABLE testDB.testschema.fooTable TO $1 AS OF SYSTEM TIME '%s'`, ts1.AsOfSystemTime()), backupTestSchemaPath) + sqlDB.Exec(t, `INSERT INTO testDB.testschema.fooTable(id) SELECT * FROM generate_series(4,10)`) + ts2 := hlc.Timestamp{WallTime: timeutil.Now().UnixNano()} + sqlDB.Exec(t, fmt.Sprintf(`BACKUP TABLE testDB.testschema.fooTable TO $1 AS OF SYSTEM TIME '%s'`, ts2.AsOfSystemTime()), backupTestSchemaPath) + testCasesOnError := []struct { name string tableName string @@ -294,6 +301,7 @@ func TestExportData(t *testing.T) { } for _, tc := range testCasesOnError { t.Run(tc.name, func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=%s --external-io-dir=%s", strings.Join(tc.backupPaths, " "), tc.tableName, @@ -308,32 +316,56 @@ func TestExportData(t *testing.T) { tableName string backupPaths []string expectedDatums string + flags string }{ { "show-data-with-qualified-table-name-of-user-defined-schema", "testDB.testschema.fooTable", []string{backupTestSchemaPath}, "2,223,'dog'\n", + "", }, { "show-data-with-qualified-table-name-of-public-schema", "testDB.public.fooTable", []string{backupPublicSchemaPath}, "1,123,'cat'\n", + "", }, { "show-data-of-incremental-backup", "testDB.testschema.fooTable", []string{backupTestSchemaPath, backupTestSchemaPath + ts1.GoTime().Format(backupccl.DateBasedIncFolderName)}, "2,223,'dog'\n3,333,'mickey mouse'\n", + "", + }, { + "show-data-of-incremental-backup-with-maxRows-flag", + "testDB.testschema.fooTable", + []string{backupTestSchemaPath, + backupTestSchemaPath + ts1.GoTime().Format(backupccl.DateBasedIncFolderName), + backupTestSchemaPath + ts2.GoTime().Format(backupccl.DateBasedIncFolderName), + }, + "2,223,'dog'\n3,333,'mickey mouse'\n4,null,null\n", + "--max-rows=3", + }, { + "show-data-of-incremental-backup-with-start-key-specified", + "testDB.testschema.fooTable", + []string{backupTestSchemaPath, + backupTestSchemaPath + ts1.GoTime().Format(backupccl.DateBasedIncFolderName), + backupTestSchemaPath + ts2.GoTime().Format(backupccl.DateBasedIncFolderName), + }, + "5,null,null\n6,null,null\n7,null,null\n8,null,null\n9,null,null\n10,null,null\n", + "--start-key=raw:\\xBF\\x89\\x8c\\x8c", }, } for _, tc := range testCasesDatumOutput { t.Run(tc.name, func(t *testing.T) { - out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=%s --external-io-dir=%s", + setDebugContextDefault() + out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=%s --external-io-dir=%s %s", strings.Join(tc.backupPaths, " "), tc.tableName, - dir)) + dir, + tc.flags)) require.NoError(t, err) checkExpectedOutput(t, tc.expectedDatums, out) }) @@ -380,6 +412,7 @@ func TestExportDataWithMultipleRanges(t *testing.T) { require.Equal(t, 8, rangeNum) t.Run("export-data-with-multiple-ranges", func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=testDB.public.fooTable --external-io-dir=%s", backupPath, dir)) @@ -392,6 +425,7 @@ func TestExportDataWithMultipleRanges(t *testing.T) { }) t.Run("export-data-with-multiple-ranges-in-incremental-backups", func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s %s --table=testDB.public.fooTable --external-io-dir=%s", backupPath, backupPath+ts.GoTime().Format(backupccl.DateBasedIncFolderName), dir)) @@ -454,6 +488,7 @@ func TestExportDataAOST(t *testing.T) { sqlDB.Exec(t, fmt.Sprintf(`BACKUP TO $1 AS OF SYSTEM TIME '%s' WITH revision_history`, ts3.AsOfSystemTime()), backupPathWithRev) t.Run("show-data-as-of-a-uncovered-timestamp", func(t *testing.T) { + setDebugContextDefault() tsNotCovered := hlc.Timestamp{WallTime: timeutil.Now().UnixNano()} out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=%s --as-of=%s --external-io-dir=%s", backupPath, @@ -468,6 +503,7 @@ func TestExportDataAOST(t *testing.T) { }) t.Run("show-data-as-of-non-backup-ts-should-return-error", func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=%s --as-of=%s --external-io-dir=%s", backupPath, "testDB.public.fooTable", @@ -631,6 +667,7 @@ func TestExportDataAOST(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { + setDebugContextDefault() out, err := c.RunWithCapture(fmt.Sprintf("debug backup export %s --table=%s --as-of=%s --external-io-dir=%s ", strings.Join(tc.backupPaths, " "), tc.tableName, diff --git a/pkg/cli/cliflags/flags.go b/pkg/cli/cliflags/flags.go index a68dcc59bd81..301adf4e7b69 100644 --- a/pkg/cli/cliflags/flags.go +++ b/pkg/cli/cliflags/flags.go @@ -1558,4 +1558,14 @@ Only csv is supported at the moment. Name: "nullas", Description: `The string that should be used to represent NULL values. `, } + + StartKey = FlagInfo{ + Name: "start-key", + Description: From.Description, + } + + MaxRows = FlagInfo{ + Name: "max-rows", + Description: `Maximum number of rows to return. `, + } ) diff --git a/pkg/cli/flags.go b/pkg/cli/flags.go index 1a448357ea26..339a53b950e0 100644 --- a/pkg/cli/flags.go +++ b/pkg/cli/flags.go @@ -825,8 +825,8 @@ func init() { // Debug commands. { f := debugKeysCmd.Flags() - varFlag(f, (*mvccKey)(&debugCtx.startKey), cliflags.From) - varFlag(f, (*mvccKey)(&debugCtx.endKey), cliflags.To) + varFlag(f, (*MVCCKey)(&debugCtx.startKey), cliflags.From) + varFlag(f, (*MVCCKey)(&debugCtx.endKey), cliflags.To) intFlag(f, &debugCtx.maxResults, cliflags.Limit) boolFlag(f, &debugCtx.values, cliflags.Values) boolFlag(f, &debugCtx.sizes, cliflags.Sizes) diff --git a/pkg/cli/flags_util.go b/pkg/cli/flags_util.go index f7ed86c5c930..660ff6adf076 100644 --- a/pkg/cli/flags_util.go +++ b/pkg/cli/flags_util.go @@ -168,18 +168,19 @@ func (m *dumpMode) Set(s string) error { return nil } -type mvccKey storage.MVCCKey +// MVCCKey implements the pflag.Value interface. +type MVCCKey storage.MVCCKey // Type implements the pflag.Value interface. -func (k *mvccKey) Type() string { return "engine.MVCCKey" } +func (k *MVCCKey) Type() string { return "engine.MVCCKey" } // String implements the pflag.Value interface. -func (k *mvccKey) String() string { +func (k *MVCCKey) String() string { return storage.MVCCKey(*k).String() } // Set implements the pflag.Value interface. -func (k *mvccKey) Set(value string) error { +func (k *MVCCKey) Set(value string) error { var typ keyType var keyStr string i := strings.IndexByte(value, ':') @@ -207,26 +208,26 @@ func (k *mvccKey) Set(value string) error { "encoded MVCCKey (i.e. with a timestamp component); here's one with a zero timestamp: %s", encoded) } - *k = mvccKey(newK) + *k = MVCCKey(newK) case raw: unquoted, err := unquoteArg(keyStr) if err != nil { return err } - *k = mvccKey(storage.MakeMVCCMetadataKey(roachpb.Key(unquoted))) + *k = MVCCKey(storage.MakeMVCCMetadataKey(roachpb.Key(unquoted))) case human: scanner := keysutil.MakePrettyScanner(nil /* tableParser */) key, err := scanner.Scan(keyStr) if err != nil { return err } - *k = mvccKey(storage.MakeMVCCMetadataKey(key)) + *k = MVCCKey(storage.MakeMVCCMetadataKey(key)) case rangeID: fromID, err := parseRangeID(keyStr) if err != nil { return err } - *k = mvccKey(storage.MakeMVCCMetadataKey(keys.MakeRangeIDPrefix(fromID))) + *k = MVCCKey(storage.MakeMVCCMetadataKey(keys.MakeRangeIDPrefix(fromID))) default: return fmt.Errorf("unknown key type %s", typ) }