Merge #141850

141850: sql/schemachanger: fix unusable secondary indexes during pk swap r=fqazi a=fqazi

Previously, the declarative schema changer during ALTER PRIMARY KEY, the declarative schema changer would make the new primary index public and then work on creating replacement secondary indexes. This could lead to broken behaviour where you could be left with either unusable old secondary indexes (causing trouble for DML queries) or potential validation errors. To address this, this patch adjusts the rules so that the new primary index is kept at a validated state, while replacement secondary indexes are constructed. This guarantees that all secondary indexes are usable during a PK swap, and the new primary key and replacement secondary indexes are swapped in an atomic fashion.

Additionally, to help support this change the index backfiller has been updated, so that a source index can be specified for backfilling indexes. This allows the index backfiller to use a validated primary index to create the secondary index, without exposing it to read queries.

Fixes: #133129
Fixes: #130165

Release note (bug fix): Addressed a bug where secondary indexes could be
unusable by DML while a primary key swap was occurring, if the new
primary key did not contain columns from the old primary key.

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

pkg/ccl/changefeedccl/schemafeed/testdata/alter_primary_key

@@ -17,29 +17,31 @@ t 2->3: Unknown
 t 3->4: Unknown
 t 4->5: Unknown
 t 5->6: Unknown
-t 6->7: PrimaryKeyChange
+t 6->7: Unknown
 t 7->8: Unknown
 t 8->9: Unknown
 t 9->10: Unknown
 t 10->11: Unknown
-t 11->12: Unknown
+t 11->12: PrimaryKeyChange
 t 12->13: Unknown
+t 13->14: Unknown
 
 exec
 ALTER TABLE t ALTER PRIMARY KEY USING COLUMNS (k, j);
 ----
 
 pop f=1
 ----
-t 13->14: Unknown
 t 14->15: Unknown
 t 15->16: Unknown
 t 16->17: Unknown
 t 17->18: Unknown
-t 18->19: PrimaryKeyChange
+t 18->19: Unknown
 t 19->20: Unknown
 t 20->21: Unknown
 t 21->22: Unknown
 t 22->23: Unknown
 t 23->24: Unknown
-t 24->25: Unknown
+t 24->25: PrimaryKeyChange
+t 25->26: Unknown
+t 26->27: Unknown

pkg/ccl/schemachangerccl/backup_base_generated_test.go

@@ -525,6 +525,7 @@ ElementState:
     isNotVisible: false
     isUnique: false
     recreateSourceId: 0
+    recreateTargetId: 0
     sharding: null
     sourceIndexId: 0
     tableId: 104
@@ -960,6 +961,7 @@ ElementState:
     isNotVisible: false
     isUnique: true
     recreateSourceId: 0
+    recreateTargetId: 0
     sharding: null
     sourceIndexId: 0
     tableId: 105

pkg/ccl/schemachangerccl/testdata/decomp/partitioning

@@ -43,7 +43,6 @@ go_test(
         "//pkg/sql",
         "//pkg/sql/catalog/lease",
         "//pkg/sql/catalog/replication",
-        "//pkg/sql/colexecerror",
         "//pkg/sql/gcjob",
         "//pkg/sql/isql",
         "//pkg/sql/lexbase",

pkg/ccl/testccl/sqlccl/BUILD.bazel

@@ -13,7 +13,6 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/internal/sqlsmith"
-	"github.com/cockroachdb/cockroach/pkg/sql/colexecerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/lexbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltestutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/tests"
@@ -139,11 +138,6 @@ func TestExplainGist(t *testing.T) {
 	skip.UnderDeadlock(t, "the test is too slow")
 	skip.UnderRace(t, "the test is too slow")
 
-	// Use the release-build panic-catching behavior instead of the
-	// crdb_test-build behavior. This is needed so that some known bugs like
-	// #119045 and #133129 don't result in a test failure.
-	defer colexecerror.ProductionBehaviorForTests()()
-
 	ctx := context.Background()
 	rng, _ := randutil.NewTestRand()
 

pkg/ccl/testccl/sqlccl/explain_test.go

@@ -1045,7 +1045,7 @@ func (sc *SchemaChanger) distIndexBackfill(
 		)
 		indexBatchSize := indexBackfillBatchSize.Get(&sc.execCfg.Settings.SV)
 		chunkSize := sc.getChunkSize(indexBatchSize)
-		spec, err := initIndexBackfillerSpec(*tableDesc.TableDesc(), writeAsOf, readAsOf, writeAtRequestTimestamp, chunkSize, addedIndexes)
+		spec, err := initIndexBackfillerSpec(*tableDesc.TableDesc(), writeAsOf, readAsOf, writeAtRequestTimestamp, chunkSize, addedIndexes, 0 /* sourceIndexID*/)
 		if err != nil {
 			return err
 		}

pkg/sql/backfill.go

@@ -26,6 +26,7 @@ go_library(
         "//pkg/sql/execinfra",
         "//pkg/sql/execinfrapb",
         "//pkg/sql/isql",
+        "//pkg/sql/parser",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/row",

pkg/sql/backfill/BUILD.bazel

@@ -480,6 +480,10 @@ type IndexBackfiller struct {
 	// backfilled.
 	indexesToEncode []catalog.Index
 
+	// sourceIndex the primary index that should be used to execute this
+	// backfill.
+	sourceIndex catalog.Index
+
 	// keyPrefixes is a slice of key prefixes for each index in indexesToEncode.
 	// indexesToEncode and keyPrefixes should both have the same ordering.
 	keyPrefixes [][]byte
@@ -513,7 +517,7 @@ func (ib *IndexBackfiller) InitForLocalUse(
 ) error {
 
 	// Initialize ib.added.
-	ib.initIndexes(evalCtx.Codec, desc, nil /* allowList */)
+	ib.initIndexes(evalCtx.Codec, desc, nil /* allowList */, 0 /*sourceIndex*/)
 
 	// Initialize ib.cols and ib.colIdxMap.
 	if err := ib.initCols(desc); err != nil {
@@ -649,14 +653,15 @@ func (ib *IndexBackfiller) InitForDistributedUse(
 	flowCtx *execinfra.FlowCtx,
 	desc catalog.TableDescriptor,
 	allowList []catid.IndexID,
+	sourceIndexID catid.IndexID,
 	mon *mon.BytesMonitor,
 ) error {
 	// We'll be modifying the eval.Context in BuildIndexEntriesChunk, so we need
 	// to make a copy.
 	evalCtx := flowCtx.NewEvalCtx()
 
 	// Initialize ib.added.
-	ib.initIndexes(evalCtx.Codec, desc, allowList)
+	ib.initIndexes(evalCtx.Codec, desc, allowList, sourceIndexID)
 
 	// Initialize ib.indexBackfillerCols.
 	if err := ib.initCols(desc); err != nil {
@@ -734,7 +739,7 @@ func (ib *IndexBackfiller) ShrinkBoundAccount(ctx context.Context, shrinkBy int6
 // populates the cols and colIdxMap fields.
 func (ib *IndexBackfiller) initCols(desc catalog.TableDescriptor) (err error) {
 	ib.indexBackfillerCols, err = makeIndexBackfillColumns(
-		desc.DeletableColumns(), desc.GetPrimaryIndex(), ib.added,
+		desc, desc.DeletableColumns(), ib.sourceIndex, ib.added,
 	)
 	return err
 }
@@ -746,7 +751,10 @@ func (ib *IndexBackfiller) initCols(desc catalog.TableDescriptor) (err error) {
 // If `allowList` is non-nil, we only add those in this list.
 // If `allowList` is nil, we add all adding index mutations.
 func (ib *IndexBackfiller) initIndexes(
-	codec keys.SQLCodec, desc catalog.TableDescriptor, allowList []catid.IndexID,
+	codec keys.SQLCodec,
+	desc catalog.TableDescriptor,
+	allowList []catid.IndexID,
+	sourceIndexID catid.IndexID,
 ) {
 	var allowListAsSet catid.IndexSet
 	if len(allowList) > 0 {
@@ -755,6 +763,11 @@ func (ib *IndexBackfiller) initIndexes(
 
 	mutations := desc.AllMutations()
 	mutationID := mutations[0].MutationID()
+	if sourceIndexID != 0 {
+		ib.sourceIndex = catalog.FindIndexByID(desc, sourceIndexID)
+	} else {
+		ib.sourceIndex = desc.GetPrimaryIndex()
+	}
 	ib.keyPrefixes = make([][]byte, 0, len(ib.added))
 	// Mutations in the same transaction have the same ID. Loop through the
 	// mutations and collect all index mutations.
@@ -860,7 +873,7 @@ func (ib *IndexBackfiller) BuildIndexEntriesChunk(
 	// read or used
 	var spec fetchpb.IndexFetchSpec
 	if err := rowenc.InitIndexFetchSpec(
-		&spec, ib.evalCtx.Codec, tableDesc, tableDesc.GetPrimaryIndex(), fetcherCols,
+		&spec, ib.evalCtx.Codec, tableDesc, ib.sourceIndex, fetcherCols,
 	); err != nil {
 		return nil, nil, memUsedPerChunk, err
 	}

pkg/sql/backfill/backfill.go

@@ -8,6 +8,8 @@ package backfill
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemaexpr"
+	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/util/intsets"
 	"github.com/cockroachdb/errors"
 )
@@ -52,7 +54,10 @@ type indexBackfillerCols struct {
 // requirements due to references in expressions. That needs to be added after
 // constructing this information.
 func makeIndexBackfillColumns(
-	deletableColumns []catalog.Column, sourcePrimaryIndex catalog.Index, addedIndexes []catalog.Index,
+	tableDesc catalog.TableDescriptor,
+	deletableColumns []catalog.Column,
+	sourcePrimaryIndex catalog.Index,
+	addedIndexes []catalog.Index,
 ) (indexBackfillerCols, error) {
 
 	// We will need to evaluate default or computed expressions for
@@ -71,6 +76,14 @@ func makeIndexBackfillColumns(
 	for _, idx := range addedIndexes {
 		allIndexColumns.UnionWith(indexColumns(idx))
 	}
+	if tableDesc != nil {
+		var err error
+		// Add any dependent columns needed to evaluate expressions.
+		allIndexColumns, err = addIndexColumnsFromExpressions(allIndexColumns, tableDesc, addedIndexes)
+		if err != nil {
+			return indexBackfillerCols{}, err
+		}
+	}
 	for _, column := range deletableColumns {
 		if !column.Public() &&
 			// Include columns we are adding, in case we are adding them to a
@@ -86,6 +99,23 @@ func makeIndexBackfillColumns(
 				allIndexColumns.Contains(column.GetID())) {
 			continue
 		}
+		// Public columns that are not in the source primary index and not
+		// needed by the secondary index can be skipped. If there are virtual
+		// columns that are needed by the secondary index then pick those up.
+		if column.Public() && !primaryColumns.Contains(column.GetID()) {
+			// If a non-virtual column requested by the secondary index is missing
+			// we are going to error out.
+			if allIndexColumns.Contains(column.GetID()) && !column.IsVirtual() {
+				return indexBackfillerCols{}, errors.AssertionFailedf(
+					"column %s is public but not in the source primary index",
+					column.GetName(),
+				)
+			}
+			// If the column is not needed by the secondary index then we can skip it.
+			if !allIndexColumns.Contains(column.GetID()) {
+				continue
+			}
+		}
 		if column.IsComputed() && column.IsVirtual() {
 			computedVirtual.Add(column.GetID())
 			ib.computedCols = append(ib.computedCols, column)
@@ -183,3 +213,45 @@ func indexColumns(idx catalog.Index) (s catalog.TableColSet) {
 	s.UnionWith(idx.CollectSecondaryStoredColumnIDs())
 	return s
 }
+
+// addIndexColumnsFromExpressions takes a set of columns stored in an index,
+// and computes any dependent columns needed for computed expressions or partial
+// indexes.
+func addIndexColumnsFromExpressions(
+	s catalog.TableColSet, table catalog.TableDescriptor, addIndexes []catalog.Index,
+) (catalog.TableColSet, error) {
+	addReferencesFromExpression := func(expression string) error {
+		expr, err := parser.ParseExpr(expression)
+		if err != nil {
+			return err
+		}
+		referencedColumns, err := schemaexpr.ExtractColumnIDs(table, expr)
+		if err != nil {
+			return err
+		}
+		for _, colID := range referencedColumns.Ordered() {
+			s.Add(colID)
+		}
+		return nil
+	}
+	// First get any columns needed to compute virtual expressions.
+	for _, colID := range s.Ordered() {
+		column := catalog.FindColumnByID(table, colID)
+		if !column.IsVirtual() || !column.IsComputed() {
+			continue
+		}
+		if err := addReferencesFromExpression(column.GetComputeExpr()); err != nil {
+			return s, err
+		}
+	}
+	// Next get any expressions needed to evaluate index predicates.
+	for _, idx := range addIndexes {
+		if len(idx.GetPredicate()) == 0 {
+			continue
+		}
+		if err := addReferencesFromExpression(idx.GetPredicate()); err != nil {
+			return s, err
+		}
+	}
+	return s, nil
+}

pkg/sql/backfill/index_backfiller_cols.go

@@ -113,9 +113,8 @@ func TestIndexBackfillerColumns(t *testing.T) {
 					keyCols: colIDs{1},
 				},
 			},
-			expCols:     colIDs{1, 2, 3},
-			expComputed: colIDs{3},
-			expNeeded:   colIDs{1},
+			expCols:   colIDs{1, 2},
+			expNeeded: colIDs{1},
 		},
 		{
 			name: "one virtual, one computed mutation column in primary",
@@ -333,7 +332,7 @@ func TestIndexBackfillerColumns(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			out, err := makeIndexBackfillColumns(
-				asColumnSlice(tc.cols), tc.src, asIndexSlice(tc.toEncode),
+				nil, asColumnSlice(tc.cols), tc.src, asIndexSlice(tc.toEncode),
 			)
 			if tc.expErr != "" {
 				require.Regexp(t, tc.expErr, err)
@@ -412,15 +411,15 @@ func TestInitIndexesAllowList(t *testing.T) {
 	t.Run("nil allowList", func(t *testing.T) {
 		// A nil allowList means no filtering.
 		ib := &IndexBackfiller{}
-		ib.initIndexes(keys.SystemSQLCodec, desc, nil /* allowList */)
+		ib.initIndexes(keys.SystemSQLCodec, desc, nil /* allowList */, 0 /* sourceIndexID */)
 		require.Equal(t, 2, len(ib.added))
 		require.Equal(t, catid.IndexID(2), ib.added[0].GetID())
 		require.Equal(t, catid.IndexID(3), ib.added[1].GetID())
 	})
 
 	t.Run("non-nil allowList", func(t *testing.T) {
 		ib := &IndexBackfiller{}
-		ib.initIndexes(keys.SystemSQLCodec, desc, []catid.IndexID{3} /* allowList */)
+		ib.initIndexes(keys.SystemSQLCodec, desc, []catid.IndexID{3} /* allowList */, 0 /* sourceIndexID */)
 		require.Equal(t, 1, len(ib.added))
 		require.Equal(t, catid.IndexID(3), ib.added[0].GetID())
 	})