diff --git a/pkg/kv/txn.go b/pkg/kv/txn.go index a18f2d971c0c..71031d58616f 100644 --- a/pkg/kv/txn.go +++ b/pkg/kv/txn.go @@ -1003,7 +1003,8 @@ func (txn *Txn) exec(ctx context.Context, fn func(context.Context, *Txn) error) // PrepareForRetry needs to be called before a retry to perform some // book-keeping and clear errors when possible. func (txn *Txn) PrepareForRetry(ctx context.Context) { - // TODO(andrei): I think commit triggers are reset in the wrong place. See #18170. + // Reset commit triggers. These must be reconfigured by the client during the + // next retry. txn.commitTriggers = nil txn.mu.Lock() diff --git a/pkg/sql/backfill/BUILD.bazel b/pkg/sql/backfill/BUILD.bazel index 005511d6b172..fc677ae58cea 100644 --- a/pkg/sql/backfill/BUILD.bazel +++ b/pkg/sql/backfill/BUILD.bazel @@ -32,6 +32,7 @@ go_library( "//pkg/sql/row", "//pkg/sql/rowenc", "//pkg/sql/rowinfra", + "//pkg/sql/sem/catid", "//pkg/sql/sem/eval", "//pkg/sql/sem/transform", "//pkg/sql/sem/tree", @@ -58,6 +59,8 @@ go_test( "//pkg/sql/catalog", "//pkg/sql/catalog/catenumpb", "//pkg/sql/catalog/descpb", + "//pkg/sql/catalog/tabledesc", + "//pkg/sql/sem/catid", "@com_github_stretchr_testify//require", ], ) diff --git a/pkg/sql/backfill/backfill.go b/pkg/sql/backfill/backfill.go index 0d8a5fdd93c1..d01b184bc014 100644 --- a/pkg/sql/backfill/backfill.go +++ b/pkg/sql/backfill/backfill.go @@ -32,6 +32,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/row" "github.com/cockroachdb/cockroach/pkg/sql/rowenc" "github.com/cockroachdb/cockroach/pkg/sql/rowinfra" + "github.com/cockroachdb/cockroach/pkg/sql/sem/catid" "github.com/cockroachdb/cockroach/pkg/sql/sem/eval" "github.com/cockroachdb/cockroach/pkg/sql/sem/transform" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" @@ -505,7 +506,7 @@ func (ib *IndexBackfiller) InitForLocalUse( ) error { // Initialize ib.added. - ib.initIndexes(desc) + ib.initIndexes(desc, nil /* allowList */) // Initialize ib.cols and ib.colIdxMap. if err := ib.initCols(desc); err != nil { @@ -640,11 +641,12 @@ func (ib *IndexBackfiller) InitForDistributedUse( ctx context.Context, flowCtx *execinfra.FlowCtx, desc catalog.TableDescriptor, + allowList []catid.IndexID, mon *mon.BytesMonitor, ) error { // Initialize ib.added. - ib.initIndexes(desc) + ib.initIndexes(desc, allowList) // Initialize ib.indexBackfillerCols. if err := ib.initCols(desc); err != nil { @@ -728,19 +730,25 @@ func (ib *IndexBackfiller) initCols(desc catalog.TableDescriptor) (err error) { } // initIndexes is a helper to populate index metadata of an IndexBackfiller. It -// populates the added field. It returns a set of column ordinals that must be -// fetched in order to backfill the added indexes. -func (ib *IndexBackfiller) initIndexes(desc catalog.TableDescriptor) { +// populates the added field to be all adding index mutations. +// 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(desc catalog.TableDescriptor, allowList []catid.IndexID) { + var allowListAsSet catid.IndexSet + if len(allowList) > 0 { + allowListAsSet = catid.MakeIndexIDSet(allowList...) + } + mutations := desc.AllMutations() mutationID := mutations[0].MutationID() - // Mutations in the same transaction have the same ID. Loop through the // mutations and collect all index mutations. for _, m := range mutations { if m.MutationID() != mutationID { break } - if IndexMutationFilter(m) { + if IndexMutationFilter(m) && + (allowListAsSet.Empty() || allowListAsSet.Contains(m.AsIndex().GetID())) { idx := m.AsIndex() ib.added = append(ib.added, idx) } diff --git a/pkg/sql/backfill/index_backfiller_cols_test.go b/pkg/sql/backfill/index_backfiller_cols_test.go index 4a2281a79f4a..d1934ea8b6eb 100644 --- a/pkg/sql/backfill/index_backfiller_cols_test.go +++ b/pkg/sql/backfill/index_backfiller_cols_test.go @@ -16,6 +16,8 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" + "github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc" + "github.com/cockroachdb/cockroach/pkg/sql/sem/catid" "github.com/stretchr/testify/require" ) @@ -367,6 +369,67 @@ func TestIndexBackfillerColumns(t *testing.T) { } } +// TestInitIndexesAllowList tests that initIndexes works correctly with +// "allowList" to populate the "added" field of the index backfiller. +func TestInitIndexesAllowList(t *testing.T) { + desc := &tabledesc.Mutable{} + desc.TableDescriptor = descpb.TableDescriptor{ + Mutations: []descpb.DescriptorMutation{ + { + // candidate 1 + Descriptor_: &descpb.DescriptorMutation_Index{ + Index: &descpb.IndexDescriptor{ID: 2}, + }, + Direction: descpb.DescriptorMutation_ADD, + }, + { + // candidate 2 + Descriptor_: &descpb.DescriptorMutation_Index{ + Index: &descpb.IndexDescriptor{ID: 3}, + }, + Direction: descpb.DescriptorMutation_ADD, + }, + { + // non-candidate: index is being dropped + Descriptor_: &descpb.DescriptorMutation_Index{ + Index: &descpb.IndexDescriptor{ID: 4}, + }, + Direction: descpb.DescriptorMutation_DROP, + }, + { + // non-candidate: index is temporary index + Descriptor_: &descpb.DescriptorMutation_Index{ + Index: &descpb.IndexDescriptor{ID: 4, UseDeletePreservingEncoding: true}, + }, + Direction: descpb.DescriptorMutation_ADD, + }, + { + // non-candidate: not an index + Descriptor_: &descpb.DescriptorMutation_Column{ + Column: &descpb.ColumnDescriptor{ID: 2}, + }, + Direction: descpb.DescriptorMutation_ADD, + }, + }, + } + + t.Run("nil allowList", func(t *testing.T) { + // A nil allowList means no filtering. + ib := &IndexBackfiller{} + ib.initIndexes(desc, nil /* allowList */) + 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(desc, []catid.IndexID{3} /* allowList */) + require.Equal(t, 1, len(ib.added)) + require.Equal(t, catid.IndexID(3), ib.added[0].GetID()) + }) +} + type fakeColumn struct { catalog.Column id descpb.ColumnID diff --git a/pkg/sql/rowexec/indexbackfiller.go b/pkg/sql/rowexec/indexbackfiller.go index 0e180d97007d..db5df157568b 100644 --- a/pkg/sql/rowexec/indexbackfiller.go +++ b/pkg/sql/rowexec/indexbackfiller.go @@ -84,7 +84,7 @@ func newIndexBackfiller( } if err := ib.IndexBackfiller.InitForDistributedUse(ctx, flowCtx, ib.desc, - indexBackfillerMon); err != nil { + ib.spec.IndexesToBackfill, indexBackfillerMon); err != nil { return nil, err } diff --git a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go index 9f7bffc042ec..26cc0856a09c 100644 --- a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go +++ b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go @@ -473,7 +473,7 @@ func getNextStoredIndexColumnOrdinal(allTargets ElementResultSet, idx *scpb.Prim func getImplicitSecondaryIndexName( b BuildCtx, descID descpb.ID, indexID descpb.IndexID, numImplicitColumns int, ) string { - elts := b.QueryByID(descID).Filter(notAbsentTargetFilter) + elts := b.QueryByID(descID).Filter(notFilter(absentTargetFilter)) var idx *scpb.Index scpb.ForEachSecondaryIndex(elts, func(current scpb.Status, target scpb.TargetStatus, e *scpb.SecondaryIndex) { if e.IndexID == indexID { diff --git a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_primary_key.go b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_primary_key.go index 72551dd341ea..2e05c0c99556 100644 --- a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_primary_key.go +++ b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_primary_key.go @@ -363,7 +363,7 @@ func fallBackIfConcurrentSchemaChange(b BuildCtx, t alterPrimaryKeySpec, tableID // fallBackIfPartitionedIndexExists panics with an unimplemented error // if there exists partitioned indexes on the table. func fallBackIfPartitionedIndexExists(b BuildCtx, t alterPrimaryKeySpec, tableID catid.DescID) { - tableElts := b.QueryByID(tableID).Filter(notAbsentTargetFilter) + tableElts := b.QueryByID(tableID).Filter(notFilter(absentTargetFilter)) scpb.ForEachIndexPartitioning(tableElts, func(_ scpb.Status, _ scpb.TargetStatus, _ *scpb.IndexPartitioning) { panic(scerrors.NotImplementedErrorf(t.n, "ALTER PRIMARY KEY on a table with index partitioning is not yet supported")) @@ -373,7 +373,7 @@ func fallBackIfPartitionedIndexExists(b BuildCtx, t alterPrimaryKeySpec, tableID // fallBackIfShardedIndexExists panics with an unimplemented // error if there exists sharded indexes on the table. func fallBackIfShardedIndexExists(b BuildCtx, t alterPrimaryKeySpec, tableID catid.DescID) { - tableElts := b.QueryByID(tableID).Filter(notAbsentTargetFilter) + tableElts := b.QueryByID(tableID).Filter(notFilter(absentTargetFilter)) var hasSecondary bool scpb.ForEachSecondaryIndex(tableElts, func(_ scpb.Status, _ scpb.TargetStatus, idx *scpb.SecondaryIndex) { hasSecondary = true diff --git a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/drop_index.go b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/drop_index.go index bb45459a37c3..b30a4d5143b1 100644 --- a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/drop_index.go +++ b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/drop_index.go @@ -506,7 +506,7 @@ func isIndexUniqueAndCanServeFK( func hasColsUniquenessConstraintOtherThan( b BuildCtx, tableID descpb.ID, columnIDs []descpb.ColumnID, otherThan descpb.ConstraintID, ) (ret bool) { - b.QueryByID(tableID).Filter(publicTargetFilter).Filter(statusPublicFilter). + b.QueryByID(tableID).Filter(publicTargetFilter).Filter(publicStatusFilter). ForEachElementStatus(func( current scpb.Status, target scpb.TargetStatus, e scpb.Element, ) { diff --git a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go index f7a807d2754e..d9e4f4d6e941 100644 --- a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go +++ b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go @@ -392,6 +392,25 @@ func isColumnFilter(_ scpb.Status, _ scpb.TargetStatus, e scpb.Element) bool { return isColumn } +func orFilter( + fs ...func(_ scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool, +) func(_ scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool { + return func(status scpb.Status, target scpb.TargetStatus, e scpb.Element) (ret bool) { + for _, f := range fs { + ret = ret || f(status, target, e) + } + return ret + } +} + +func notFilter( + f func(_ scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool, +) func(_ scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool { + return func(status scpb.Status, target scpb.TargetStatus, e scpb.Element) bool { + return !f(status, target, e) + } +} + func publicTargetFilter(_ scpb.Status, target scpb.TargetStatus, _ scpb.Element) bool { return target == scpb.ToPublic } @@ -400,18 +419,16 @@ func absentTargetFilter(_ scpb.Status, target scpb.TargetStatus, _ scpb.Element) return target == scpb.ToAbsent } -func notAbsentTargetFilter(_ scpb.Status, target scpb.TargetStatus, _ scpb.Element) bool { - return target != scpb.ToAbsent +func publicStatusFilter(status scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool { + return status == scpb.Status_PUBLIC } -func statusAbsentOrBackfillOnlyFilter( - status scpb.Status, _ scpb.TargetStatus, _ scpb.Element, -) bool { - return status == scpb.Status_ABSENT || status == scpb.Status_BACKFILL_ONLY +func absentStatusFilter(status scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool { + return status == scpb.Status_ABSENT } -func statusPublicFilter(status scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool { - return status == scpb.Status_PUBLIC +func backfillOnlyStatusFilter(status scpb.Status, _ scpb.TargetStatus, _ scpb.Element) bool { + return status == scpb.Status_BACKFILL_ONLY } func hasIndexIDAttrFilter( @@ -470,8 +487,8 @@ func getPrimaryIndexes( allTargets := b.QueryByID(tableID) _, _, freshlyAdded = scpb.FindPrimaryIndex(allTargets. Filter(publicTargetFilter). - Filter(statusAbsentOrBackfillOnlyFilter)) - _, _, existing = scpb.FindPrimaryIndex(allTargets.Filter(statusPublicFilter)) + Filter(orFilter(absentStatusFilter, backfillOnlyStatusFilter))) + _, _, existing = scpb.FindPrimaryIndex(allTargets.Filter(publicStatusFilter)) if existing == nil { // TODO(postamar): can this even be possible? panic(pgerror.Newf(pgcode.NoPrimaryKey, "missing active primary key")) @@ -578,7 +595,7 @@ func (s indexSpec) clone() (c indexSpec) { // makeIndexSpec constructs an indexSpec based on an existing index element. func makeIndexSpec(b BuildCtx, tableID catid.DescID, indexID catid.IndexID) (s indexSpec) { - tableElts := b.QueryByID(tableID).Filter(notAbsentTargetFilter) + tableElts := b.QueryByID(tableID).Filter(notFilter(absentTargetFilter)) idxElts := tableElts.Filter(hasIndexIDAttrFilter(indexID)) var constraintID catid.ConstraintID var n int @@ -718,7 +735,7 @@ func makeSwapIndexSpec( var inID, tempID catid.IndexID var inConstraintID catid.ConstraintID { - _, _, tbl := scpb.FindTable(b.QueryByID(tableID).Filter(notAbsentTargetFilter)) + _, _, tbl := scpb.FindTable(b.QueryByID(tableID).Filter(notFilter(absentTargetFilter))) inID = b.NextTableIndexID(tbl) inConstraintID = b.NextTableConstraintID(tbl.TableID) tempID = inID + 1 diff --git a/pkg/sql/sem/tree/create.go b/pkg/sql/sem/tree/create.go index 6adda8f5c1c2..ae440b751d6e 100644 --- a/pkg/sql/sem/tree/create.go +++ b/pkg/sql/sem/tree/create.go @@ -1161,6 +1161,11 @@ func (node *UniqueConstraintTableDef) Format(ctx *FmtCtx) { ctx.WriteString(" VISIBILITY " + fmt.Sprintf("%.2f", 1-invisibility)) } } + if node.StorageParams != nil { + ctx.WriteString(" WITH (") + ctx.FormatNode(&node.StorageParams) + ctx.WriteString(")") + } } // ForeignKeyConstraintTableDef represents a FOREIGN KEY constraint in the AST. diff --git a/pkg/sql/sem/tree/pretty.go b/pkg/sql/sem/tree/pretty.go index 9325c1522f3a..d3d2b255a1cb 100644 --- a/pkg/sql/sem/tree/pretty.go +++ b/pkg/sql/sem/tree/pretty.go @@ -1826,6 +1826,13 @@ func (node *UniqueConstraintTableDef) doc(p *PrettyCfg) pretty.Doc { } } + if node.StorageParams != nil { + clauses = append(clauses, p.bracketKeyword( + "WITH", "(", + p.Doc(&node.StorageParams), + ")", "")) + } + if len(clauses) == 0 { return title }