From dd32aa8382cf47fb7763991bc1cef104cd6da55a Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Wed, 10 Mar 2021 12:20:50 -0500 Subject: [PATCH] sql: stop mutating AST in AlterPrimaryKey The code in #61345 added a test that, under stress, nicely revealed a bug. The bug is that we're mutating the AST in place and then, on retries, we hit an error. The fix is to not mutate the AST. I wish I had a more comprehensive testing strategy to ensure this didn't happen. On some level, we could clone the AST when passing it to various DDL plan node constructors. That's very defensive but also probably fine. Another thing that would be cool would be to just assert that after planning the AST did not change. Touches #60824. Release note: None --- pkg/sql/alter_primary_key.go | 9 +++++---- pkg/sql/alter_table.go | 4 ++-- pkg/sql/alter_table_locality.go | 2 +- pkg/sql/create_index.go | 35 +++++++++++++++++++-------------- pkg/sql/create_table.go | 34 ++++++++++++++++++++------------ 5 files changed, 49 insertions(+), 35 deletions(-) diff --git a/pkg/sql/alter_primary_key.go b/pkg/sql/alter_primary_key.go index 635d2107f553..1b79694ffc69 100644 --- a/pkg/sql/alter_primary_key.go +++ b/pkg/sql/alter_primary_key.go @@ -48,7 +48,7 @@ type alterPrimaryKeyLocalitySwap struct { func (p *planner) AlterPrimaryKey( ctx context.Context, tableDesc *tabledesc.Mutable, - alterPKNode *tree.AlterTableAlterPrimaryKey, + alterPKNode tree.AlterTableAlterPrimaryKey, alterPrimaryKeyLocalitySwap *alterPrimaryKeyLocalitySwap, ) error { if alterPKNode.Interleave != nil { @@ -153,7 +153,7 @@ func (p *planner) AlterPrimaryKey( // primary index, which would mean nothing needs to be modified // here. { - requiresIndexChange, err := p.shouldCreateIndexes(ctx, tableDesc, alterPKNode, alterPrimaryKeyLocalitySwap) + requiresIndexChange, err := p.shouldCreateIndexes(ctx, tableDesc, &alterPKNode, alterPrimaryKeyLocalitySwap) if err != nil { return err } @@ -190,12 +190,12 @@ func (p *planner) AlterPrimaryKey( // If the new index is requested to be sharded, set up the index descriptor // to be sharded, and add the new shard column if it is missing. if alterPKNode.Sharded != nil { - shardCol, newColumn, err := setupShardedIndex( + shardCol, newColumns, newColumn, err := setupShardedIndex( ctx, p.EvalContext(), &p.semaCtx, p.SessionData().HashShardedIndexesEnabled, - &alterPKNode.Columns, + alterPKNode.Columns, alterPKNode.Sharded.ShardBuckets, tableDesc, newPrimaryIndexDesc, @@ -204,6 +204,7 @@ func (p *planner) AlterPrimaryKey( if err != nil { return err } + alterPKNode.Columns = newColumns if newColumn { if err := p.setupFamilyAndConstraintForShard( ctx, diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go index 7987ebd6e74f..db419db4419e 100644 --- a/pkg/sql/alter_table.go +++ b/pkg/sql/alter_table.go @@ -220,7 +220,7 @@ func (n *alterTableNode) startExec(params runParams) error { if err := params.p.AlterPrimaryKey( params.ctx, n.tableDesc, - alterPK, + *alterPK, nil, /* localityConfigSwap */ ); err != nil { return err @@ -373,7 +373,7 @@ func (n *alterTableNode) startExec(params runParams) error { if err := params.p.AlterPrimaryKey( params.ctx, n.tableDesc, - t, + *t, nil, /* localityConfigSwap */ ); err != nil { return err diff --git a/pkg/sql/alter_table_locality.go b/pkg/sql/alter_table_locality.go index 8d53f8f96067..a13aea66e903 100644 --- a/pkg/sql/alter_table_locality.go +++ b/pkg/sql/alter_table_locality.go @@ -430,7 +430,7 @@ func (n *alterTableSetLocalityNode) alterTableLocalityFromOrToRegionalByRow( if err := params.p.AlterPrimaryKey( params.ctx, n.tableDesc, - &tree.AlterTableAlterPrimaryKey{ + tree.AlterTableAlterPrimaryKey{ Name: tree.Name(n.tableDesc.PrimaryIndex.Name), Columns: cols, }, diff --git a/pkg/sql/create_index.go b/pkg/sql/create_index.go index 191aebc65a44..5e3bf7e04b65 100644 --- a/pkg/sql/create_index.go +++ b/pkg/sql/create_index.go @@ -145,7 +145,7 @@ func (p *planner) setupFamilyAndConstraintForShard( // is hash sharded. Note that `tableDesc` will be modified when this method is called for // a hash sharded index. func MakeIndexDescriptor( - params runParams, n *tree.CreateIndex, tableDesc *tabledesc.Mutable, + params runParams, n tree.CreateIndex, tableDesc *tabledesc.Mutable, ) (*descpb.IndexDescriptor, error) { // Ensure that the columns we want to index exist before trying to create the // index. @@ -200,7 +200,7 @@ func MakeIndexDescriptor( } telemetry.Inc(sqltelemetry.InvertedIndexCounter) } - + columns := n.Columns if n.Sharded != nil { if n.PartitionByIndex.ContainsPartitions() { return nil, pgerror.New(pgcode.FeatureNotSupported, "sharded indexes don't support partitioning") @@ -211,12 +211,12 @@ func MakeIndexDescriptor( if n.Interleave != nil { return nil, pgerror.New(pgcode.FeatureNotSupported, "interleaved indexes cannot also be hash sharded") } - shardCol, newColumn, err := setupShardedIndex( + shardCol, newColumns, newColumn, err := setupShardedIndex( params.ctx, params.EvalContext(), ¶ms.p.semaCtx, params.SessionData().HashShardedIndexesEnabled, - &n.Columns, + n.Columns, n.Sharded.ShardBuckets, tableDesc, &indexDesc, @@ -224,6 +224,7 @@ func MakeIndexDescriptor( if err != nil { return nil, err } + columns = newColumns if newColumn { if err := params.p.setupFamilyAndConstraintForShard(params.ctx, tableDesc, shardCol, indexDesc.Sharded.ColumnNames, indexDesc.Sharded.ShardBuckets); err != nil { @@ -243,7 +244,7 @@ func MakeIndexDescriptor( telemetry.Inc(sqltelemetry.PartialIndexCounter) } - if err := indexDesc.FillColumns(n.Columns); err != nil { + if err := indexDesc.FillColumns(columns); err != nil { return nil, err } @@ -285,46 +286,50 @@ func (n *createIndexNode) ReadingOwnWrites() {} var hashShardedIndexesDisabledError = pgerror.Newf(pgcode.FeatureNotSupported, "hash sharded indexes require the experimental_enable_hash_sharded_indexes session variable") +// setupShardedIndex updates the index descriptor with the relevant new column. +// It also returns the column so it can be added to the table. It returns the +// new column set for the index. This set must be used regardless of whether or +// not there is a newly created column. func setupShardedIndex( ctx context.Context, evalCtx *tree.EvalContext, semaCtx *tree.SemaContext, shardedIndexEnabled bool, - columns *tree.IndexElemList, + columns tree.IndexElemList, bucketsExpr tree.Expr, tableDesc *tabledesc.Mutable, indexDesc *descpb.IndexDescriptor, isNewTable bool, -) (shard *descpb.ColumnDescriptor, newColumn bool, err error) { +) (shard *descpb.ColumnDescriptor, newColumns tree.IndexElemList, newColumn bool, err error) { if !shardedIndexEnabled { - return nil, false, hashShardedIndexesDisabledError + return nil, nil, false, hashShardedIndexesDisabledError } - colNames := make([]string, 0, len(*columns)) - for _, c := range *columns { + colNames := make([]string, 0, len(columns)) + for _, c := range columns { colNames = append(colNames, string(c.Column)) } buckets, err := tabledesc.EvalShardBucketCount(ctx, semaCtx, evalCtx, bucketsExpr) if err != nil { - return nil, false, err + return nil, nil, false, err } shardCol, newColumn, err := maybeCreateAndAddShardCol(int(buckets), tableDesc, colNames, isNewTable) if err != nil { - return nil, false, err + return nil, nil, false, err } shardIdxElem := tree.IndexElem{ Column: tree.Name(shardCol.Name), Direction: tree.Ascending, } - *columns = append(tree.IndexElemList{shardIdxElem}, *columns...) + newColumns = append(tree.IndexElemList{shardIdxElem}, columns...) indexDesc.Sharded = descpb.ShardedDescriptor{ IsSharded: true, Name: shardCol.Name, ShardBuckets: buckets, ColumnNames: colNames, } - return shardCol, newColumn, nil + return shardCol, newColumns, newColumn, nil } // maybeCreateAndAddShardCol adds a new hidden computed shard column (or its mutation) to @@ -437,7 +442,7 @@ func (n *createIndexNode) startExec(params runParams) error { } } - indexDesc, err := MakeIndexDescriptor(params, n.n, n.tableDesc) + indexDesc, err := MakeIndexDescriptor(params, *n.n, n.tableDesc) if err != nil { return err } diff --git a/pkg/sql/create_table.go b/pkg/sql/create_table.go index e0657de829e5..817cd538bfad 100644 --- a/pkg/sql/create_table.go +++ b/pkg/sql/create_table.go @@ -1877,36 +1877,38 @@ func NewTableDesc( } } - setupShardedIndexForNewTable := func(d *tree.IndexTableDef, idx *descpb.IndexDescriptor) error { + setupShardedIndexForNewTable := func( + d tree.IndexTableDef, idx *descpb.IndexDescriptor, + ) (columns tree.IndexElemList, _ error) { if n.PartitionByTable.ContainsPartitions() { - return pgerror.New(pgcode.FeatureNotSupported, "sharded indexes don't support partitioning") + return nil, pgerror.New(pgcode.FeatureNotSupported, "sharded indexes don't support partitioning") } - shardCol, newColumn, err := setupShardedIndex( + shardCol, newColumns, newColumn, err := setupShardedIndex( ctx, evalCtx, semaCtx, sessionData.HashShardedIndexesEnabled, - &d.Columns, + d.Columns, d.Sharded.ShardBuckets, &desc, idx, true /* isNewTable */) if err != nil { - return err + return nil, err } if newColumn { buckets, err := tabledesc.EvalShardBucketCount(ctx, semaCtx, evalCtx, d.Sharded.ShardBuckets) if err != nil { - return err + return nil, err } checkConstraint, err := makeShardCheckConstraintDef(&desc, int(buckets), shardCol) if err != nil { - return err + return nil, err } n.Defs = append(n.Defs, checkConstraint) columnDefaultExprs = append(columnDefaultExprs, nil) } - return nil + return newColumns, nil } idxValidator := schemaexpr.MakeIndexPredicateValidator(ctx, n.Table, &desc, semaCtx) @@ -1929,6 +1931,7 @@ func NewTableDesc( if d.Inverted { idx.Type = descpb.IndexDescriptor_INVERTED } + columns := d.Columns if d.Sharded != nil { if d.Interleave != nil { return nil, pgerror.New(pgcode.FeatureNotSupported, "interleaved indexes cannot also be hash sharded") @@ -1936,11 +1939,13 @@ func NewTableDesc( if isRegionalByRow { return nil, hashShardedIndexesOnRegionalByRowError() } - if err := setupShardedIndexForNewTable(d, &idx); err != nil { + var err error + columns, err = setupShardedIndexForNewTable(*d, &idx) + if err != nil { return nil, err } } - if err := idx.FillColumns(d.Columns); err != nil { + if err := idx.FillColumns(columns); err != nil { return nil, err } if d.Inverted { @@ -2024,6 +2029,7 @@ func NewTableDesc( StoreColumnNames: d.Storing.ToStrings(), Version: indexEncodingVersion, } + columns := d.Columns if d.Sharded != nil { if n.Interleave != nil && d.PrimaryKey { return nil, pgerror.New(pgcode.FeatureNotSupported, "interleaved indexes cannot also be hash sharded") @@ -2031,11 +2037,13 @@ func NewTableDesc( if isRegionalByRow { return nil, hashShardedIndexesOnRegionalByRowError() } - if err := setupShardedIndexForNewTable(&d.IndexTableDef, &idx); err != nil { + var err error + columns, err = setupShardedIndexForNewTable(d.IndexTableDef, &idx) + if err != nil { return nil, err } } - if err := idx.FillColumns(d.Columns); err != nil { + if err := idx.FillColumns(columns); err != nil { return nil, err } // Specifying a partitioning on a PRIMARY KEY constraint should be disallowed by the @@ -2096,7 +2104,7 @@ func NewTableDesc( "interleave not supported in primary key constraint definition", ) } - for _, c := range d.Columns { + for _, c := range columns { primaryIndexColumnSet[string(c.Column)] = struct{}{} } }