sql: add validation for CTAS / materialized views

Previously, no validation existed when backfilling queries for CTAS
statements or creating materialized view. This could lead to scenarios
with unnoticed data loss if anything went wrong. To address this, this
patch adds validation operations that backfill queries into tables to
confirm that row counts match the source.

Fixes: #144957

Release note: None

Author: fqazi

pkg/sql/schema_changer.go

@@ -308,6 +308,88 @@ func (sc *SchemaChanger) refreshMaterializedView(
 
 const schemaChangerBackfillTxnDebugName = "schemaChangerBackfill"
 
+// validateBackfillQueryIntoTable validates that source query matches the contents of
+// a backfilled table, when executing the query at queryTS.
+func (sc *SchemaChanger) validateBackfillQueryIntoTable(
+	ctx context.Context, table catalog.TableDescriptor, queryTS hlc.Timestamp, query string,
+) error {
+	var entryCount int64
+	sd := NewInternalSessionData(ctx, sc.execCfg.Settings, "validateBackfillQueryIntoTable")
+	sd.SessionData = *sc.sessionData
+	// First get the expected row count for the source query at the target timestamp.
+	if err := sc.execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		parsedQuery, err := parser.ParseOne(query)
+		if err != nil {
+			return err
+		}
+		// If the query has an AOST clause, then we will remove it here.
+		selectTop, ok := parsedQuery.AST.(*tree.Select)
+		if ok {
+			selectStmt := selectTop.Select
+			var parenSel *tree.ParenSelect
+			var ok bool
+			for parenSel, ok = selectStmt.(*tree.ParenSelect); ok; parenSel, ok = selectStmt.(*tree.ParenSelect) {
+				selectStmt = parenSel.Select.Select
+			}
+			sc, ok := selectStmt.(*tree.SelectClause)
+			if ok && sc.From.AsOf.Expr != nil {
+				sc.From.AsOf.Expr = nil
+				query = parsedQuery.AST.String()
+			}
+		}
+		// Inject the query and the time we should scan at.
+		err = txn.KV().SetFixedTimestamp(ctx, queryTS)
+		if err != nil {
+			return err
+		}
+		countQuery := fmt.Sprintf("SELECT count(*) FROM (%s)", query)
+		row, err := txn.QueryRow(ctx, "backfill-query-src-count", txn.KV(), countQuery)
+		if err != nil {
+			return err
+		}
+		entryCount = int64(tree.MustBeDInt(row[0]))
+		return nil
+	}, isql.WithSessionData(sd)); err != nil {
+		return err
+	}
+	// Next run validation on table that was populated using count queries.
+	// Get rid of the table ID prefix.
+	index := table.GetPrimaryIndex()
+	now := sc.db.KV().Clock().Now()
+	builder := table.NewBuilder()
+	// Make the table public so that we can validate our expected
+	// counts match.
+	mut := builder.BuildExistingMutable().(*tabledesc.Mutable)
+	mut.SetPublic()
+	count, err := countIndexRowsAndMaybeCheckUniqueness(ctx, mut, index, false,
+		descs.NewHistoricalInternalExecTxnRunner(now, func(ctx context.Context, fn descs.InternalExecFn) error {
+			return sc.execCfg.InternalDB.DescsTxn(ctx, func(
+				ctx context.Context, txn descs.Txn,
+			) error {
+				if err := txn.KV().SetFixedTimestamp(ctx, now); err != nil {
+					return err
+				}
+				return fn(ctx, txn)
+			}, isql.WithPriority(admissionpb.BulkNormalPri))
+		}),
+		sessiondata.NodeUserWithBulkLowPriSessionDataOverride)
+	if err != nil {
+		return err
+	}
+	// Testing knob that allows us to manipulate counts to fail
+	// the validation.
+	if sc.testingKnobs.RunDuringQueryBackfillValidation != nil {
+		count, err = sc.testingKnobs.RunDuringQueryBackfillValidation(entryCount, count)
+		if err != nil {
+			return err
+		}
+	}
+	if count != entryCount {
+		return errors.AssertionFailedf("backfill query did not populate index %q with expected number of rows (expected: %d, got: %d)", index.GetName(), entryCount, count)
+	}
+	return nil
+}
+
 func (sc *SchemaChanger) backfillQueryIntoTable(
 	ctx context.Context,
 	table catalog.TableDescriptor,
@@ -334,6 +416,9 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 			}
 		}
 	}()
+	validationTime := ts
+	var skipValidation bool
+	// Get the expected entry count against the source query.
 	err = sc.db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
 		defer func() {
 			isTxnRetry = true
@@ -430,7 +515,6 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 				return err
 			}
 		}
-
 		res := kvpb.BulkOpSummary{}
 		rw := NewCallbackResultWriter(func(ctx context.Context, row tree.Datums) error {
 			// TODO(adityamaru): Use the BulkOpSummary for either telemetry or to
@@ -477,7 +561,6 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 					}
 				}
 			}
-
 			planDistribution, _ := getPlanDistribution(
 				ctx, localPlanner.Descriptors().HasUncommittedTypes(),
 				localPlanner.extendedEvalCtx.SessionData(),
@@ -495,7 +578,22 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 			}
 		})
 
-		return planAndRunErr
+		if planAndRunErr != nil {
+			return planAndRunErr
+		}
+		// Otherwise, the select statement had no fixed timestamp. For
+		// validating counts we will start with the read timestamp.
+		if ts.IsEmpty() {
+			validationTime = txn.KV().ReadTimestamp()
+		}
+		// If a virtual table was used then we can't conduct any kind of validation,
+		// since our AOST query might be reading data not in KV.
+		for _, tbl := range localPlanner.curPlan.mem.Metadata().AllTables() {
+			if tbl.Table.IsVirtualTable() {
+				skipValidation = true
+			}
+		}
+		return nil
 	})
 
 	// BatchTimestampBeforeGCError is retryable for the schema changer, but we
@@ -506,7 +604,18 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 			errors.Wrap(err, "unable to retry backfill since fixed timestamp is before the GC timestamp"),
 		)
 	}
-	return err
+	if err != nil {
+		return err
+	}
+
+	// Validation will be skipped if the query is not AOST safe.
+	// (i.e. reading non-KV data via CRDB internal)
+	if !skipValidation {
+		if err := sc.validateBackfillQueryIntoTable(ctx, table, validationTime, query); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 // maybe backfill a created table by executing the AS query. Return nil if
@@ -2773,6 +2882,10 @@ type SchemaChangerTestingKnobs struct {
 
 	// RunBeforeModifyRowLevelTTL is called just before the modify row level TTL is committed.
 	RunBeforeModifyRowLevelTTL func() error
+
+	// RunAfterQueryBackfillValidation is called right after validation is executed
+	// for a query backfill.
+	RunDuringQueryBackfillValidation func(expectedCount int64, currentCount int64) (newCurrentCount int64, err error)
 }
 
 // ModuleTestingKnobs is part of the base.ModuleTestingKnobs interface.

pkg/sql/schema_changer_test.go

@@ -7883,3 +7883,31 @@ func TestLeaseGenerationBumpWithSchemaChange(t *testing.T) {
 	runner.Exec(t, "ALTER TABLE t1 ALTER PRIMARY KEY USING COLUMNS(n, j)")
 	require.NoError(t, grp.Wait())
 }
+
+// TestCreateTableAsValidationFailure simulates a synthetic validation
+// failure for CREATE TABLE AS.
+func TestCreateTableAsValidationFailure(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			SQLSchemaChanger: &sql.SchemaChangerTestingKnobs{
+				RunDuringQueryBackfillValidation: func(expectedCount int64, currentCount int64) (newCurrentCount int64, err error) {
+					return 0, nil
+				},
+			},
+		},
+	})
+
+	defer s.Stopper().Stop(context.Background())
+	runner := sqlutils.MakeSQLRunner(sqlDB)
+	// Create table table and populate it.
+	runner.Exec(t, "CREATE TABLE t1(n int PRIMARY KEY)")
+	runner.Exec(t, "INSERT INTO t1 VALUES (1)")
+	runner.Exec(t, "INSERT INTO t1 VALUES (2)")
+	runner.Exec(t, "INSERT INTO t1 VALUES (3)")
+	// Execute a CTAS and CREATE MATERIALIZED VIEW statements that should fail.
+	runner.ExpectErr(t, "backfill query did not populate index \"t2_pkey\" with expected number of rows", "CREATE TABLE t2 AS (SELECT * FROM t1)")
+	runner.ExpectErr(t, "backfill query did not populate index \"t2_pkey\" with expected number of rows", "CREATE MATERIALIZED VIEW t2 AS (SELECT n FROM t1)")
+}