opt: fix race caused by shared table annotations

pkg/sql/opt/column_meta.go

@@ -15,9 +15,6 @@
 package opt
 
 import (
-	"strings"
-
-	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/types"
 	"github.com/cockroachdb/cockroach/pkg/util"
 )
@@ -58,73 +55,9 @@ type ColumnMeta struct {
 	// Type is the scalar SQL type of this column.
 	Type types.T
 
-	// TableMeta is the metadata for the base table to which this column belongs.
-	// If the column was synthesized (i.e. no base table), then is is null.
-	TableMeta *TableMeta
-
-	// md is a back-reference to the query metadata.
-	md *Metadata
-}
-
-// QualifiedAlias returns the column alias, possibly qualified with the table,
-// schema, or database name:
-//
-//   1. If fullyQualify is true, then the returned alias is prefixed by the
-//      original, fully qualified name of the table: tab.Name().FQString().
-//
-//   2. If there's another column in the metadata with the same column alias but
-//      a different table name, then prefix the column alias with the table
-//      name: "tabName.columnAlias".
-//
-func (cm *ColumnMeta) QualifiedAlias(fullyQualify bool) string {
-	if cm.TableMeta == nil {
-		// Column doesn't belong to a table, so no need to qualify it further.
-		return cm.Alias
-	}
-	md := cm.md
-
-	// If a fully qualified alias has not been requested, then only qualify it if
-	// it would otherwise be ambiguous.
-	var tabAlias tree.TableName
-	qualify := fullyQualify
-	if !fullyQualify {
-		for i := range md.cols {
-			if i == int(cm.MetaID-1) {
-				continue
-			}
-
-			// If there are two columns with same alias, then column is ambiguous.
-			cm2 := &md.cols[i]
-			if cm2.Alias == cm.Alias {
-				tabAlias = cm.TableMeta.Alias
-				if cm2.TableMeta == nil {
-					qualify = true
-				} else {
-					// Only qualify if the qualified names are actually different.
-					tabAlias2 := cm2.TableMeta.Alias
-					if tabAlias.String() != tabAlias2.String() {
-						qualify = true
-					}
-				}
-			}
-		}
-	}
-
-	// If the column name should not even be partly qualified, then no more to do.
-	if !qualify {
-		return cm.Alias
-	}
-
-	var sb strings.Builder
-	if fullyQualify {
-		s := cm.TableMeta.Table.Name().FQString()
-		sb.WriteString(s)
-	} else {
-		sb.WriteString(tabAlias.String())
-	}
-	sb.WriteRune('.')
-	sb.WriteString(cm.Alias)
-	return sb.String()
+	// Table is the base table to which this column belongs.
+	// If the column was synthesized (i.e. no base table), then it is 0.
+	Table TableID
 }
 
 // ToSet converts a column id list to a column id set.

pkg/sql/opt/exec/execbuilder/format.go

@@ -64,7 +64,7 @@ func fmtInterceptor(f *memo.ExprFmtCtx, tp treeprinter.Node, nd opt.Expr) bool {
 	fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
 	fmtCtx.SetIndexedVarFormat(func(ctx *tree.FmtCtx, idx int) {
 		fullyQualify := !f.HasFlags(memo.ExprFmtHideQualifications)
-		alias := md.ColumnMeta(opt.ColumnID(idx + 1)).QualifiedAlias(fullyQualify)
+		alias := md.QualifiedAlias(opt.ColumnID(idx+1), fullyQualify)
 		ctx.WriteString(alias)
 	})
 	expr.Format(fmtCtx)

pkg/sql/opt/memo/expr_format.go

@@ -715,7 +715,7 @@ func formatCol(
 	colMeta := md.ColumnMeta(id)
 	if label == "" {
 		fullyQualify := !f.HasFlags(ExprFmtHideQualifications)
-		label = colMeta.QualifiedAlias(fullyQualify)
+		label = md.QualifiedAlias(id, fullyQualify)
 	}
 
 	if !isSimpleColumnName(label) {
@@ -752,7 +752,7 @@ func FormatPrivate(f *ExprFmtCtx, private interface{}, physProps *physical.Requi
 	switch t := private.(type) {
 	case *opt.ColumnID:
 		fullyQualify := !f.HasFlags(ExprFmtHideQualifications)
-		label := f.Memo.metadata.ColumnMeta(*t).QualifiedAlias(fullyQualify)
+		label := f.Memo.metadata.QualifiedAlias(*t, fullyQualify)
 		fmt.Fprintf(f.Buffer, " %s", label)
 
 	case *TupleOrdinal:

pkg/sql/opt/memo/memo_test.go

@@ -21,11 +21,10 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
-	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
+	opttestutils "github.com/cockroachdb/cockroach/pkg/sql/opt/testutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/opttester"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testcat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
-	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
@@ -61,24 +60,10 @@ func TestMemoInit(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	stmt, err := parser.ParseOne("SELECT * FROM abc WHERE $1=10")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	ctx := context.Background()
-	semaCtx := tree.MakeSemaContext()
-	if err := semaCtx.Placeholders.Init(stmt.NumPlaceholders, nil /* typeHints */); err != nil {
-		t.Fatal(err)
-	}
 	evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
 
 	var o xform.Optimizer
-	o.Init(&evalCtx)
-	err = optbuilder.New(ctx, &semaCtx, &evalCtx, catalog, o.Factory(), stmt.AST).Build()
-	if err != nil {
-		t.Fatal(err)
-	}
+	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT * FROM abc WHERE $1=10")
 
 	o.Init(&evalCtx)
 	if !o.Memo().IsEmpty() {
@@ -110,27 +95,16 @@ func TestMemoIsStale(t *testing.T) {
 	// access via the view.
 	catalog.Table(tree.NewTableName("t", "abc")).Revoked = true
 
-	ctx := context.Background()
-	semaCtx := tree.MakeSemaContext()
-	evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
-
 	// Initialize context with starting values.
+	evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
 	evalCtx.SessionData.Database = "t"
 
-	stmt, err := parser.ParseOne("SELECT a, b+1 FROM abcview WHERE c='foo'")
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	var o xform.Optimizer
-	o.Init(&evalCtx)
-	err = optbuilder.New(ctx, &semaCtx, &evalCtx, catalog, o.Factory(), stmt.AST).Build()
-	if err != nil {
-		t.Fatal(err)
-	}
+	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT a, b+1 FROM abcview WHERE c='foo'")
 	o.Memo().Metadata().AddSchemaDependency(catalog.Schema().Name(), catalog.Schema(), privilege.CREATE)
 	o.Memo().Metadata().AddSchema(catalog.Schema())
 
+	ctx := context.Background()
 	stale := func() {
 		t.Helper()
 		if isStale, err := o.Memo().IsStale(ctx, &evalCtx, catalog); err != nil {

pkg/sql/opt/metadata.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"math/bits"
+	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
@@ -88,6 +89,8 @@ type Metadata struct {
 	// might resolve to the same object now but to different objects later; we
 	// want to verify the resolution of both names.
 	deps []mdDep
+
+	// NOTE! When adding fields here, update CopyFrom.
 }
 
 type mdDep struct {
@@ -131,13 +134,25 @@ func (md *Metadata) Init() {
 
 // CopyFrom initializes the metadata with a copy of the provided metadata.
 // This metadata can then be modified independent of the copied metadata.
+//
+// Table annotations are not transferred over; all annotations are unset on
+// the copy.
 func (md *Metadata) CopyFrom(from *Metadata) {
-	if len(md.cols) != 0 || len(md.tables) != 0 || len(md.deps) != 0 {
+	if len(md.schemas) != 0 || len(md.cols) != 0 || len(md.tables) != 0 ||
+		len(md.sequences) != 0 || len(md.deps) != 0 {
 		panic("CopyFrom requires empty destination")
 	}
 	md.schemas = append(md.schemas, from.schemas...)
 	md.cols = append(md.cols, from.cols...)
 	md.tables = append(md.tables, from.tables...)
+
+	// Clear table annotations. These objects can be mutable and can't be safely
+	// shared between different metadata instances.
+	for i := range md.tables {
+		md.tables[i].clearAnnotations()
+	}
+
+	md.sequences = append(md.sequences, from.sequences...)
 	md.deps = append(md.deps, from.deps...)
 }
 
@@ -274,7 +289,6 @@ func (md *Metadata) AddTableWithAlias(tab cat.Table, alias *tree.TableName) Tabl
 		md.tables = make([]TableMeta, 0, 4)
 	}
 	md.tables = append(md.tables, TableMeta{MetaID: tabID, Table: tab, Alias: *alias})
-	tabMeta := md.TableMeta(tabID)
 
 	colCount := tab.DeletableColumnCount()
 	if md.cols == nil {
@@ -284,7 +298,7 @@ func (md *Metadata) AddTableWithAlias(tab cat.Table, alias *tree.TableName) Tabl
 	for i := 0; i < colCount; i++ {
 		col := tab.Column(i)
 		colID := md.AddColumn(string(col.ColName()), col.DatumType())
-		md.ColumnMeta(colID).TableMeta = tabMeta
+		md.ColumnMeta(colID).Table = tabID
 	}
 
 	return tabID
@@ -327,7 +341,7 @@ func (md *Metadata) AddColumn(alias string, typ types.T) ColumnID {
 		alias = fmt.Sprintf("column%d", len(md.cols)+1)
 	}
 	colID := ColumnID(len(md.cols) + 1)
-	md.cols = append(md.cols, ColumnMeta{MetaID: colID, Alias: alias, Type: typ, md: md})
+	md.cols = append(md.cols, ColumnMeta{MetaID: colID, Alias: alias, Type: typ})
 	return colID
 }
 
@@ -343,6 +357,67 @@ func (md *Metadata) ColumnMeta(colID ColumnID) *ColumnMeta {
 	return &md.cols[colID.index()]
 }
 
+// QualifiedAlias returns the column alias, possibly qualified with the table,
+// schema, or database name:
+//
+//   1. If fullyQualify is true, then the returned alias is prefixed by the
+//      original, fully qualified name of the table: tab.Name().FQString().
+//
+//   2. If there's another column in the metadata with the same column alias but
+//      a different table name, then prefix the column alias with the table
+//      name: "tabName.columnAlias".
+//
+func (md *Metadata) QualifiedAlias(colID ColumnID, fullyQualify bool) string {
+	cm := md.ColumnMeta(colID)
+	if cm.Table == 0 {
+		// Column doesn't belong to a table, so no need to qualify it further.
+		return cm.Alias
+	}
+
+	// If a fully qualified alias has not been requested, then only qualify it if
+	// it would otherwise be ambiguous.
+	var tabAlias tree.TableName
+	qualify := fullyQualify
+	if !fullyQualify {
+		for i := range md.cols {
+			if i == int(cm.MetaID-1) {
+				continue
+			}
+
+			// If there are two columns with same alias, then column is ambiguous.
+			cm2 := &md.cols[i]
+			if cm2.Alias == cm.Alias {
+				tabAlias = md.TableMeta(cm.Table).Alias
+				if cm2.Table == 0 {
+					qualify = true
+				} else {
+					// Only qualify if the qualified names are actually different.
+					tabAlias2 := md.TableMeta(cm2.Table).Alias
+					if tabAlias.String() != tabAlias2.String() {
+						qualify = true
+					}
+				}
+			}
+		}
+	}
+
+	// If the column name should not even be partly qualified, then no more to do.
+	if !qualify {
+		return cm.Alias
+	}
+
+	var sb strings.Builder
+	if fullyQualify {
+		s := md.TableMeta(cm.Table).Table.Name().FQString()
+		sb.WriteString(s)
+	} else {
+		sb.WriteString(tabAlias.String())
+	}
+	sb.WriteRune('.')
+	sb.WriteString(cm.Alias)
+	return sb.String()
+}
+
 // SequenceID uniquely identifies the usage of a sequence within the scope of a
 // query. SequenceID 0 is reserved to mean "unknown sequence".
 type SequenceID uint64

pkg/sql/opt/norm/factory.go

@@ -165,19 +165,25 @@ func (f *Factory) CustomFuncs() *CustomFuncs {
 // The "replace" callback function allows the caller to override the default
 // traversal and cloning behavior with custom logic. It is called for each node
 // in the "from" subtree, and has the choice of constructing an arbitrary
-// replacement node, or else delegating to the default behavior by returning the
-// unchanged "e" expression. The default behavior simply constructs a copy of
-// the source operator using children returned by recursive calls to the replace
-// callback. Here is example usage:
+// replacement node, or delegating to the default behavior by calling
+// CopyAndReplaceDefault, which constructs a copy of the source operator using
+// children returned by recursive calls to the replace callback. Note that if a
+// non-leaf replacement node is constructed, its inputs must be copied using
+// CopyAndReplaceDefault.
 //
-//   f.CopyAndReplace(from, fromProps, func(e opt.Expr) opt.Expr {
+// Sample usage:
+//
+//   var replaceFn ReplaceFunc
+//   replaceFn = func(e opt.Expr) opt.Expr {
 //     if e.Op() == opt.PlaceholderOp {
 //       return f.ConstructConst(evalPlaceholder(e))
 //     }
 //
-//     // Return unchanged "e" expression to get default behavior.
-//     return e
-//   })
+//     // Copy e, calling replaceFn on its inputs recursively.
+//     return f.CopyAndReplaceDefault(e, replaceFn)
+//   }
+//
+//   f.CopyAndReplace(from, fromProps, replaceFn)
 //
 // NOTE: Callers must take care to always create brand new copies of non-
 // singleton source nodes rather than referencing existing nodes. The source
@@ -222,16 +228,18 @@ func (f *Factory) AssignPlaceholders(from *memo.Memo) (err error) {
 
 	// Copy the "from" memo to this memo, replacing any Placeholder operators as
 	// the copy proceeds.
-	f.CopyAndReplace(from.RootExpr().(memo.RelExpr), from.RootProps(), func(e opt.Expr) opt.Expr {
+	var replaceFn ReplaceFunc
+	replaceFn = func(e opt.Expr) opt.Expr {
 		if placeholder, ok := e.(*memo.PlaceholderExpr); ok {
 			d, err := e.(*memo.PlaceholderExpr).Value.Eval(f.evalCtx)
 			if err != nil {
 				panic(placeholderError{err})
 			}
 			return f.ConstructConstVal(d, placeholder.DataType())
 		}
-		return e
-	})
+		return f.CopyAndReplaceDefault(e, replaceFn)
+	}
+	f.CopyAndReplace(from.RootExpr().(memo.RelExpr), from.RootProps(), replaceFn)
 
 	return nil
 }