fingerprintutils: select database and tables with aost if specified

Previously, if a user wanted to fingerprint a database(s) with an AOST, the
utility mistakenly selected the database/tables without an aost. This patch
fixes this.

Fixes #150548

Release note: none

Author: msbutler

pkg/testutils/fingerprintutils/fingerprint.go

@@ -22,6 +22,13 @@ type FingerprintOption struct {
 	StartTime       hlc.Timestamp
 }
 
+func (f FingerprintOption) AOSTCmd() string {
+	if f.AOST.IsEmpty() {
+		return ""
+	}
+	return fmt.Sprintf("AS OF SYSTEM TIME '%s'", f.AOST.AsOfSystemTime())
+}
+
 func AOST(aost hlc.Timestamp) func(*FingerprintOption) {
 	return func(opt *FingerprintOption) {
 		opt.AOST = aost
@@ -69,10 +76,7 @@ func FingerprintTable(
 		return 0, err
 	}
 
-	aostCmd := ""
-	if !opts.AOST.IsEmpty() {
-		aostCmd = fmt.Sprintf("AS OF SYSTEM TIME '%s'", opts.AOST.AsOfSystemTime())
-	}
+	aostCmd := opts.AOSTCmd()
 
 	cmd := fmt.Sprintf(`SELECT * FROM crdb_internal.fingerprint(crdb_internal.table_span(%d),true) %s`, tableID, aostCmd)
 
@@ -115,12 +119,18 @@ func FingerprintTables(
 func FingerprintDatabase(
 	ctx context.Context, db *gosql.DB, dbName string, optFuncs ...func(*FingerprintOption),
 ) (map[string]int64, error) {
+	opts, err := getOpts(optFuncs...)
+	if err != nil {
+		return nil, err
+	}
+	aostCmd := opts.AOSTCmd()
+
 	var databaseID int
-	row := db.QueryRowContext(ctx, `SELECT id FROM system.namespace WHERE name = $1`, dbName)
+	row := db.QueryRowContext(ctx, fmt.Sprintf(`SELECT id FROM system.namespace %s WHERE name = $1`, aostCmd), dbName)
 	if err := row.Scan(&databaseID); err != nil {
 		return nil, errors.New("could not get database descriptor id")
 	}
-	rows, err := db.QueryContext(ctx, `SELECT id, name FROM system.namespace where "parentID" = $1`, databaseID)
+	rows, err := db.QueryContext(ctx, fmt.Sprintf(`SELECT id, name FROM system.namespace %s where "parentID" = $1`, aostCmd), databaseID)
 	if err != nil {
 		return nil, errors.New("could not get database table name and tableIDs")
 	}
@@ -158,9 +168,12 @@ func FingerprintDatabase(
 func FingerprintAllDatabases(
 	ctx context.Context, db *gosql.DB, includeSystemDB bool, optFuncs ...func(*FingerprintOption),
 ) (map[string]map[string]int64, error) {
-
+	opts, err := getOpts(optFuncs...)
+	if err != nil {
+		return nil, err
+	}
 	dbNames := make([]string, 0)
-	rows, err := db.QueryContext(ctx, `SELECT database_name FROM [SHOW DATABASES]`)
+	rows, err := db.QueryContext(ctx, fmt.Sprintf(`select name from system.namespace %s where "parentID"=0`, opts.AOSTCmd()))
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get database names")
 	}

pkg/testutils/fingerprintutils/fingerprint_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/stretchr/testify/require"
 )
@@ -73,3 +74,42 @@ func TestFingerprintUtility(t *testing.T) {
 	_, err = fingerprintutils.FingerprintDatabase(ctx, db, "d1", fingerprintutils.Stripped(), fingerprintutils.RevisionHistory())
 	require.ErrorContains(t, err, "cannot specify stripped and revision history")
 }
+
+// TestFingerprintAllDatabasesWithAOST tests that FingerprintAllDatabases works
+// correctly when using AOST (As Of System Time) even when the target
+// database/tables are deleted after the AOST timestamp.
+func TestFingerprintAllDatabasesWithAOST(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	ctx := context.Background()
+	tc := testcluster.StartTestCluster(t, 1, base.TestClusterArgs{})
+	defer tc.Stopper().Stop(ctx)
+	db := tc.ServerConn(0)
+	sql := sqlutils.MakeSQLRunner(db)
+
+	rng, _ := randutil.NewTestRand()
+
+	sql.Exec(t, `CREATE DATABASE test_aost`)
+	sql.Exec(t, `CREATE TABLE test_aost.table_to_delete (id INT PRIMARY KEY, data STRING)`)
+	sql.Exec(t, `INSERT INTO test_aost.table_to_delete VALUES (1, 'test_data')`)
+
+	aostTimestamp := hlc.Timestamp{WallTime: timeutil.Now().UnixNano()}
+
+	sql.Exec(t, `DROP TABLE test_aost.table_to_delete`)
+	if rng.Intn(2) == 0 {
+		sql.Exec(t, `DROP DATABASE test_aost`)
+	}
+
+	// FingerprintAllDatabases with AOST should still work and find the deleted table
+	fingerprintsAOST, err := fingerprintutils.FingerprintAllDatabases(ctx, db, false,
+		fingerprintutils.AOST(aostTimestamp))
+	require.NoError(t, err)
+
+	testAostFingerprints, exists := fingerprintsAOST["test_aost"]
+	require.True(t, exists, "test_aost database should be found in fingerprints")
+
+	// Verify that the deleted table is found in the AOST fingerprint
+	tableFingerprint, exists := testAostFingerprints["table_to_delete"]
+	require.True(t, exists, "table_to_delete should be found in AOST fingerprints")
+	require.NotZero(t, tableFingerprint, "fingerprint should be non-zero")
+}