backup: fix wrong URIs/localities on skip compaction chains

The current tests for `ValidateEndTimeAndTruncate` only check that the
correct number of URIs and localities are returned, but do not check
that the correct values are returned. This allowed a bug to slip by
where if compactions were being skipped, the wrong URIs and localities
would be returned. This patch fixes the bug and updates the tests to be
more specific.

Epic: None

Release note: None

Author: kev-cao

pkg/backup/backupinfo/BUILD.bazel

@@ -38,7 +38,6 @@ go_library(
         "//pkg/sql/sem/tree",
         "//pkg/sql/stats",
         "//pkg/storage",
-        "//pkg/util",
         "//pkg/util/bulk",
         "//pkg/util/ctxgroup",
         "//pkg/util/encoding",
@@ -81,6 +80,7 @@ go_test(
         "//pkg/sql/isql",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/testcluster",
+        "//pkg/util",
         "//pkg/util/bulk",
         "//pkg/util/hlc",
         "//pkg/util/leaktest",

pkg/backup/backupinfo/manifest_handling.go

@@ -42,7 +42,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
-	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/bulk"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
@@ -890,11 +889,8 @@ func ValidateEndTimeAndTruncate(
 	includeCompacted bool,
 ) ([]string, []backuppb.BackupManifest, []jobspb.RestoreDetails_BackupLocalityInfo, error) {
 	if !includeCompacted {
-		mainBackupManifests = util.Filter(
-			mainBackupManifests,
-			func(manifest backuppb.BackupManifest) bool {
-				return !manifest.IsCompacted
-			},
+		defaultURIs, mainBackupManifests, localityInfo = skipCompactedBackups(
+			defaultURIs, mainBackupManifests, localityInfo,
 		)
 	}
 
@@ -965,6 +961,25 @@ func ValidateEndTimeAndTruncate(
 	)
 }
 
+// skipCompactedBackups removes any compacted backups from the list of
+// backups and returns the updated list of URIs, manifests, and locality info.
+//
+// NOTE: This function modifies the underlying memory of the slices passed in.
+func skipCompactedBackups(
+	defaultURIs []string,
+	manifests []backuppb.BackupManifest,
+	localityInfo []jobspb.RestoreDetails_BackupLocalityInfo,
+) ([]string, []backuppb.BackupManifest, []jobspb.RestoreDetails_BackupLocalityInfo) {
+	for i := len(manifests) - 1; i >= 0; i-- {
+		if manifests[i].IsCompacted {
+			defaultURIs = slices.Delete(defaultURIs, i, i+1)
+			manifests = slices.Delete(manifests, i, i+1)
+			localityInfo = slices.Delete(localityInfo, i, i+1)
+		}
+	}
+	return defaultURIs, manifests, localityInfo
+}
+
 // validateContinuity checks that the backups are continuous.
 func validateContinuity(manifests []backuppb.BackupManifest) error {
 	if len(manifests) == 0 {

pkg/backup/backupinfo/manifest_handling_test.go

@@ -8,7 +8,9 @@ package backupinfo_test
 import (
 	"context"
 	"fmt"
+	"slices"
 	"sort"
+	"strconv"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/backup/backupinfo"
@@ -24,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/bulk"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
@@ -549,10 +552,32 @@ func TestValidateEndTimeAndTruncate(t *testing.T) {
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
+			// Using the expected start and end times, create a slice containing the
+			// expected expectedOrder of the manifests based on their original indexes. For
+			// example, [2, 0, 1] means that the expected output should should have
+			// the third input manifest first, the first input manifest second, and the
+			// second input manifest last.
+			expectedOrder := util.Map(tc.expected, func(ts []int) string {
+				return strconv.Itoa(slices.IndexFunc(tc.manifests, func(m backuppb.BackupManifest) bool {
+					return m.StartTime.WallTime == int64(ts[0]) && m.EndTime.WallTime == int64(ts[1])
+				}))
+			})
+			// Create URIs and locality info in a way that we can track their
+			// resulting order after being truncated. We set input URIs to their
+			// stringified index, and locality info maps to a map containing just
+			// their index. We can later use the expectedOrder to check that the
+			// output maps to the expected indexes.
+			inputURIs := make([]string, len(tc.manifests))
+			inputLocs := make([]jobspb.RestoreDetails_BackupLocalityInfo, len(tc.manifests))
+			for i := range tc.manifests {
+				index := strconv.Itoa(i)
+				inputURIs[i] = index
+				inputLocs[i].URIsByOriginalLocalityKV = make(map[string]string)
+				inputLocs[i].URIsByOriginalLocalityKV[index] = index
+			}
+
 			uris, res, locs, err := backupinfo.ValidateEndTimeAndTruncate(
-				make([]string, len(tc.manifests)),
-				tc.manifests,
-				make([]jobspb.RestoreDetails_BackupLocalityInfo, len(tc.manifests)),
+				inputURIs, tc.manifests, inputLocs,
 				hlc.Timestamp{WallTime: int64(tc.endTime)},
 				false, /* includeSkipped */
 				tc.includeCompacted,
@@ -561,9 +586,14 @@ func TestValidateEndTimeAndTruncate(t *testing.T) {
 				require.ErrorContains(t, err, tc.err)
 				return
 			}
-			require.Equal(t, len(tc.expected), len(uris))
-			require.Equal(t, len(tc.expected), len(locs))
 			require.Equal(t, len(tc.expected), len(res))
+			require.Equal(t, expectedOrder, uris)
+			require.Len(t, locs, len(tc.expected))
+			for idx, rank := range expectedOrder {
+				_, ok := locs[idx].URIsByOriginalLocalityKV[rank]
+				require.True(t, ok)
+			}
+
 			for i := range tc.expected {
 				actual := []int{int(res[i].StartTime.WallTime), int(res[i].EndTime.WallTime)}
 				require.Equal(t, tc.expected[i], actual)

pkg/backup/compaction_test.go

@@ -699,6 +699,48 @@ func TestCompactionCheckpointing(t *testing.T) {
 	validateCompactedBackupForTables(t, db, []string{"bank"}, "'nodelocal://1/backup'", start, end, 2)
 }
 
+func TestToggleCompactionForRestore(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	tc, db, _, cleanup := backupRestoreTestSetup(
+		t, singleNode, 1 /* numAccounts */, InitManualReplication,
+	)
+	defer cleanup()
+	start := getTime(t)
+	db.Exec(t, fmt.Sprintf("BACKUP INTO 'nodelocal://1/backup' AS OF SYSTEM TIME '%d'", start))
+	db.Exec(t, "BACKUP INTO LATEST IN 'nodelocal://1/backup'")
+	end := getTime(t)
+	db.Exec(t, fmt.Sprintf("BACKUP INTO LATEST IN 'nodelocal://1/backup' AS OF SYSTEM TIME '%d'", end))
+	var backupPath string
+	db.QueryRow(t, "SHOW BACKUPS IN 'nodelocal://1/backup'").Scan(&backupPath)
+	var compactionID jobspb.JobID
+	db.QueryRow(
+		t,
+		fmt.Sprintf(
+			`SELECT crdb_internal.backup_compaction(
+				'BACKUP INTO LATEST IN ''nodelocal://1/backup''',
+				'%s', %d::DECIMAL, %d::DECIMAL
+			)`,
+			backupPath, start, end,
+		),
+	).Scan(&compactionID)
+	waitForSuccessfulJob(t, tc, compactionID)
+
+	var compRestoreID, classicRestoreID jobspb.JobID
+	var unused any
+	db.QueryRow(
+		t, "RESTORE DATABASE data FROM LATEST IN 'nodelocal://1/backup' WITH new_db_name = 'data1'",
+	).Scan(&compRestoreID, &unused, &unused, &unused)
+	db.Exec(t, "SET CLUSTER SETTING restore.compacted_backups.enabled = false")
+	db.QueryRow(
+		t, "RESTORE DATABASE data FROM LATEST IN 'nodelocal://1/backup' WITH new_db_name = 'data2'",
+	).Scan(&classicRestoreID, &unused, &unused, &unused)
+
+	require.Equal(t, 2, getNumBackupsInRestore(t, db, compRestoreID))
+	require.Equal(t, 3, getNumBackupsInRestore(t, db, classicRestoreID))
+}
+
 // validateCompactedBackupForTables is a wrapper around
 // validateCompactedBackupForTablesWithOpts that passes in no options.
 func validateCompactedBackupForTables(
@@ -747,7 +789,13 @@ func validateCompactedBackupForTablesWithOpts(
 		restoredRows := db.QueryStr(t, "SELECT * FROM "+table)
 		require.Equal(t, originalRows, restoredRows, "table %s", table)
 	}
-	// Check that the number of backups used in the restore is correct.
+
+	require.Equal(t, numBackups, getNumBackupsInRestore(t, db, restoreJobID))
+}
+
+// getNumBackupsInRestore returns the number of backups used in the restore
+func getNumBackupsInRestore(t *testing.T, db *sqlutils.SQLRunner, jobID jobspb.JobID) int {
+	t.Helper()
 	var detailsStr string
 	db.QueryRow(t, `SELECT crdb_internal.pb_to_json(
 		'cockroach.sql.jobs.jobspb.Payload',
@@ -756,10 +804,10 @@ func validateCompactedBackupForTablesWithOpts(
 	FROM system.job_info 
 	WHERE job_id = $1
 	AND info_key = 'legacy_payload';
-	`, restoreJobID).Scan(&detailsStr)
+	`, jobID).Scan(&detailsStr)
 	var details jobspb.RestoreDetails
 	require.NoError(t, json.Unmarshal([]byte(detailsStr), &details))
-	require.Equal(t, numBackups, len(details.URIs))
+	return len(details.URIs)
 }
 
 // ensureBackupExists ensures that a backup exists that spans the given start and end times.