sql: Refactors CTAS logic to allow user defined PKs.

This change removes the need for a separate `AsColumnNames` field
in the CreateTable node, and uses `ColumnTableDefs` instead. This
is similar to a normal CREATE TABLE query, thereby allowing us to
piggyback on most of the existing logic. It ensures `row_id` key
generation only occurs if no PK is explicitly specified.
Some special handling is required in pretty printing, as column
types are not populated at the time of parsing.

Release note: None

Author: adityamaru27

pkg/sql/create_table.go

@@ -69,7 +69,12 @@ func (p *planner) CreateTable(ctx context.Context, n *tree.CreateTable) (planNod
 			return nil, err
 		}
 
-		numColNames := len(n.AsColumnNames)
+		numColNames := 0
+		for i := 0; i < len(n.Defs); i++ {
+			if _, ok := n.Defs[i].(*tree.ColumnTableDef); ok {
+				numColNames++
+			}
+		}
 		numColumns := len(planColumns(sourcePlan))
 		if numColNames != 0 && numColNames != numColumns {
 			sourcePlan.Close(ctx)
@@ -80,12 +85,22 @@ func (p *planner) CreateTable(ctx context.Context, n *tree.CreateTable) (planNod
 		}
 
 		// Synthesize an input column that provides the default value for the
-		// hidden rowid column.
+		// hidden rowid column, if none of the provided columns are specified
+		// as the PRIMARY KEY.
 		synthRowID = true
+		for _, def := range n.Defs {
+			if d, ok := def.(*tree.ColumnTableDef); ok && d.PrimaryKey {
+				synthRowID = false
+				break
+			}
+		}
 	}
 
 	ct := &createTableNode{n: n, dbDesc: dbDesc, sourcePlan: sourcePlan}
 	ct.run.synthRowID = synthRowID
+	// This method is only invoked if the heuristic planner was used in the
+	// planning stage.
+	ct.run.fromHeuristicPlanner = true
 	return ct, nil
 }
 
@@ -95,10 +110,17 @@ type createTableRun struct {
 	autoCommit autoCommitOpt
 
 	// synthRowID indicates whether an input column needs to be synthesized to
-	// provide the default value for the hidden rowid column. The optimizer's
-	// plan already includes this column (so synthRowID is false), whereas the
-	// heuristic planner's plan does not (so synthRowID is true).
+	// provide the default value for the hidden rowid column. The optimizer's plan
+	// already includes this column if a user specified PK does not exist (so
+	// synthRowID is false), whereas the heuristic planner's plan does not in this
+	// case (so synthRowID is true).
 	synthRowID bool
+
+	// fromHeuristicPlanner indicates whether the planning was performed by the
+	// heuristic planner instead of the optimizer. This is used to determine
+	// whether or not a row_id was synthesized as part of the planning stage, if a
+	// user defined PK is not specified.
+	fromHeuristicPlanner bool
 }
 
 func (n *createTableNode) startExec(params runParams) error {
@@ -139,15 +161,18 @@ func (n *createTableNode) startExec(params runParams) error {
 		}
 
 		asCols = planColumns(n.sourcePlan)
-		if !n.run.synthRowID {
-			// rowID column is already present in the input as the last column, so
-			// ignore it for the purpose of creating column metadata (because
+		if !n.run.fromHeuristicPlanner && !n.n.AsHasUserSpecifiedPrimaryKey() {
+			// rowID column is already present in the input as the last column if it
+			// was planned by the optimizer and the user did not specify a PRIMARY
+			// KEY. So ignore it for the purpose of creating column metadata (because
 			// makeTableDescIfAs does it automatically).
 			asCols = asCols[:len(asCols)-1]
 		}
-		desc, err = makeTableDescIfAs(
+
+		desc, err = makeTableDescIfAs(params,
 			n.n, n.dbDesc.ID, id, creationTime, asCols,
-			privs, &params.p.semaCtx, params.p.EvalContext())
+			privs, params.p.EvalContext())
+
 		if err != nil {
 			return err
 		}
@@ -259,9 +284,9 @@ func (n *createTableNode) startExec(params runParams) error {
 			return err
 		}
 
-		// Prepare the buffer for row values. At this point, one more
-		// column has been added by ensurePrimaryKey() to the list of
-		// columns in sourcePlan.
+		// Prepare the buffer for row values. At this point, one more column has
+		// been added by ensurePrimaryKey() to the list of columns in sourcePlan, if
+		// a PRIMARY KEY is not specified by the user.
 		rowBuffer := make(tree.Datums, len(desc.Columns))
 		pkColIdx := len(desc.Columns) - 1
 
@@ -975,38 +1000,51 @@ func getFinalSourceQuery(source *tree.Select, evalCtx *tree.EvalContext) string
 // makeTableDescIfAs is the MakeTableDesc method for when we have a table
 // that is created with the CREATE AS format.
 func makeTableDescIfAs(
+	params runParams,
 	p *tree.CreateTable,
 	parentID, id sqlbase.ID,
 	creationTime hlc.Timestamp,
 	resultColumns []sqlbase.ResultColumn,
 	privileges *sqlbase.PrivilegeDescriptor,
-	semaCtx *tree.SemaContext,
 	evalContext *tree.EvalContext,
 ) (desc sqlbase.MutableTableDescriptor, err error) {
-	desc = InitTableDescriptor(id, parentID, p.Table.Table(), creationTime, privileges)
-	desc.CreateQuery = getFinalSourceQuery(p.AsSource, evalContext)
-
-	for i, colRes := range resultColumns {
-		columnTableDef := tree.ColumnTableDef{Name: tree.Name(colRes.Name), Type: colRes.Typ}
-		columnTableDef.Nullable.Nullability = tree.SilentNull
-		if len(p.AsColumnNames) > i {
-			columnTableDef.Name = p.AsColumnNames[i]
-		}
-
-		// The new types in the CREATE TABLE AS column specs never use
-		// SERIAL so we need not process SERIAL types here.
-		col, _, _, err := sqlbase.MakeColumnDefDescs(&columnTableDef, semaCtx)
-		if err != nil {
-			return desc, err
+	colResIndex := 0
+	// TableDefs for a CREATE TABLE ... AS AST node comprise of a ColumnTableDef
+	// for each column, and a ConstraintTableDef for any constraints on those
+	// columns.
+	for _, defs := range p.Defs {
+		var d *tree.ColumnTableDef
+		var ok bool
+		if d, ok = defs.(*tree.ColumnTableDef); ok {
+			d.Type = resultColumns[colResIndex].Typ
+			colResIndex++
+		}
+	}
+
+	// If there are no TableDefs defined by the parser, then we construct a
+	// ColumnTableDef for each column using resultColumns.
+	if len(p.Defs) == 0 {
+		for _, colRes := range resultColumns {
+			var d *tree.ColumnTableDef
+			var ok bool
+			var tableDef tree.TableDef = &tree.ColumnTableDef{Name: tree.Name(colRes.Name), Type: colRes.Typ}
+			if d, ok = tableDef.(*tree.ColumnTableDef); !ok {
+				return desc, errors.Errorf("failed to cast type to ColumnTableDef\n")
+			}
+			d.Nullable.Nullability = tree.SilentNull
+			p.Defs = append(p.Defs, tableDef)
 		}
-		desc.AddColumn(col)
 	}
 
-	// AllocateIDs mutates its receiver. `return desc, desc.AllocateIDs()`
-	// happens to work in gc, but does not work in gccgo.
-	//
-	// See https://github.com/golang/go/issues/23188.
-	err = desc.AllocateIDs()
+	desc, err = makeTableDesc(
+		params,
+		p,
+		parentID, id,
+		creationTime,
+		privileges,
+		nil, /* affected */
+	)
+	desc.CreateQuery = getFinalSourceQuery(p.AsSource, evalContext)
 	return desc, err
 }
 

pkg/sql/logictest/testdata/logic_test/create_as

@@ -144,3 +144,95 @@ SELECT * FROM baz
 ----
 a  b  c
 1  2  4
+
+# Check that CREATE TABLE AS allows users to specify primary key (#20940)
+statement ok
+CREATE TABLE foo5 (a, b PRIMARY KEY, c) AS SELECT * FROM baz
+
+query TT
+SHOW CREATE TABLE foo5
+----
+foo5  CREATE TABLE foo5 (
+    a INT8 NULL,
+    b INT8 NOT NULL,
+    c INT8 NULL,
+    CONSTRAINT "primary" PRIMARY KEY (b ASC),
+    FAMILY "primary" (a, b, c)
+    )
+
+statement ok
+SET OPTIMIZER=ON; CREATE TABLE foo6 (a PRIMARY KEY, b, c) AS SELECT * FROM baz; SET OPTIMIZER=OFF
+
+query TT
+SHOW CREATE TABLE foo6
+----
+foo6  CREATE TABLE foo6 (
+    a INT8 NOT NULL,
+    b INT8 NULL,
+    c INT8 NULL,
+    CONSTRAINT "primary" PRIMARY KEY (a ASC),
+    FAMILY "primary" (a, b, c)
+    )
+
+statement error generate insert row: null value in column "x" violates not-null constraint
+CREATE TABLE foo7 (x PRIMARY KEY) AS VALUES (1), (NULL);
+
+statement ok
+BEGIN; CREATE TABLE foo8 (item PRIMARY KEY, qty) AS SELECT * FROM stock UNION VALUES ('spoons', 25), ('knives', 50); END
+
+query TT
+SHOW CREATE TABLE foo8
+----
+foo8  CREATE TABLE foo8 (
+    item STRING NOT NULL,
+    qty INT8 NULL,
+    CONSTRAINT "primary" PRIMARY KEY (item ASC),
+    FAMILY "primary" (item, qty)
+    )
+
+# Allow CREATE TABLE AS to specify composite primary keys.
+statement ok
+CREATE TABLE foo9 (a, b, c, PRIMARY KEY (a, c)) AS SELECT * FROM baz
+
+query TT
+SHOW CREATE TABLE foo9
+----
+foo9  CREATE TABLE foo9 (
+    a INT8 NOT NULL,
+    b INT8 NULL,
+    c INT8 NOT NULL,
+    CONSTRAINT "primary" PRIMARY KEY (a ASC, c ASC),
+    FAMILY "primary" (a, b, c)
+    )
+
+statement ok
+CREATE TABLE foo10 (a, PRIMARY KEY (c, b, a), b, c) AS SELECT * FROM foo9
+
+query TT
+SHOW CREATE TABLE foo10
+----
+foo10  CREATE TABLE foo10 (
+    a INT8 NOT NULL,
+    b INT8 NOT NULL,
+    c INT8 NOT NULL,
+    CONSTRAINT "primary" PRIMARY KEY (c ASC, b ASC, a ASC),
+    FAMILY "primary" (a, b, c)
+    )
+
+statement ok
+CREATE TABLE foo11 (x, y, z, PRIMARY KEY(x, z)) AS VALUES (1, 3, 4), (10, 20, 40);
+
+query TT
+SHOW CREATE TABLE foo11
+----
+foo11  CREATE TABLE foo11 (
+    x INT8 NOT NULL,
+    y INT8 NULL,
+    z INT8 NOT NULL,
+    CONSTRAINT "primary" PRIMARY KEY (x ASC, z ASC),
+    FAMILY "primary" (x, y, z)
+    )
+
+statement error pq: multiple primary keys for table "foo12" are not allowed
+CREATE TABLE foo12 (x PRIMARY KEY, y, PRIMARY KEY(y)) AS VALUES (1, 2), (3, 4);
+

pkg/sql/opt/optbuilder/create_table.go

@@ -41,7 +41,12 @@ func (b *Builder) buildCreateTable(ct *tree.CreateTable, inScope *scope) (outSco
 		// Build the input query.
 		outScope := b.buildSelect(ct.AsSource, nil /* desiredTypes */, inScope)
 
-		numColNames := len(ct.AsColumnNames)
+		numColNames := 0
+		for i := 0; i < len(ct.Defs); i++ {
+			if _, ok := ct.Defs[i].(*tree.ColumnTableDef); ok {
+				numColNames++
+			}
+		}
 		numColumns := len(outScope.cols)
 		if numColNames != 0 && numColNames != numColumns {
 			panic(builderError{sqlbase.NewSyntaxError(fmt.Sprintf(
@@ -50,17 +55,20 @@ func (b *Builder) buildCreateTable(ct *tree.CreateTable, inScope *scope) (outSco
 				numColumns, util.Pluralize(int64(numColumns))))})
 		}
 
-		// Synthesize rowid column, and append to end of column list.
-		props, overloads := builtins.GetBuiltinProperties("unique_rowid")
-		private := &memo.FunctionPrivate{
-			Name:       "unique_rowid",
-			Typ:        types.Int,
-			Properties: props,
-			Overload:   &overloads[0],
+		input = outScope.expr
+		if !ct.AsHasUserSpecifiedPrimaryKey() {
+			// Synthesize rowid column, and append to end of column list.
+			props, overloads := builtins.GetBuiltinProperties("unique_rowid")
+			private := &memo.FunctionPrivate{
+				Name:       "unique_rowid",
+				Typ:        types.Int,
+				Properties: props,
+				Overload:   &overloads[0],
+			}
+			fn := b.factory.ConstructFunction(memo.EmptyScalarListExpr, private)
+			scopeCol := b.synthesizeColumn(outScope, "rowid", types.Int, nil /* expr */, fn)
+			input = b.factory.CustomFuncs().ProjectExtraCol(outScope.expr, fn, scopeCol.id)
 		}
-		fn := b.factory.ConstructFunction(memo.EmptyScalarListExpr, private)
-		scopeCol := b.synthesizeColumn(outScope, "rowid", types.Int, nil /* expr */, fn)
-		input = b.factory.CustomFuncs().ProjectExtraCol(outScope.expr, fn, scopeCol.id)
 		inputCols = outScope.makePhysicalProps().Presentation
 	} else {
 		// Create dummy empty input.

pkg/sql/sem/tree/create.go

@@ -356,8 +356,14 @@ func (node *ColumnTableDef) HasColumnFamily() bool {
 // Format implements the NodeFormatter interface.
 func (node *ColumnTableDef) Format(ctx *FmtCtx) {
 	ctx.FormatNode(&node.Name)
-	ctx.WriteByte(' ')
-	ctx.WriteString(node.columnTypeString())
+
+	// ColumnTableDef node type will not be specified if it represents a CREATE
+	// TABLE ... AS query.
+	if node.Type != nil {
+		ctx.WriteByte(' ')
+		ctx.WriteString(node.columnTypeString())
+	}
+
 	if node.Nullable.Nullability != SilentNull && node.Nullable.ConstraintName != "" {
 		ctx.WriteString(" CONSTRAINT ")
 		ctx.FormatNode(&node.Nullable.ConstraintName)
@@ -885,13 +891,15 @@ func (node *RangePartition) Format(ctx *FmtCtx) {
 
 // CreateTable represents a CREATE TABLE statement.
 type CreateTable struct {
-	IfNotExists   bool
-	Table         TableName
-	Interleave    *InterleaveDef
-	PartitionBy   *PartitionBy
-	Defs          TableDefs
-	AsSource      *Select
-	AsColumnNames NameList // Only to be used in conjunction with AsSource
+	IfNotExists bool
+	Table       TableName
+	Interleave  *InterleaveDef
+	PartitionBy *PartitionBy
+	// In CREATE...AS queries, Defs represents a list of ColumnTableDefs, one for
+	// each column, and a ConstraintTableDef for each constraint on a subset of
+	// these columns.
+	Defs     TableDefs
+	AsSource *Select
 }
 
 // As returns true if this table represents a CREATE TABLE ... AS statement,
@@ -900,6 +908,21 @@ func (node *CreateTable) As() bool {
 	return node.AsSource != nil
 }
 
+// AsHasUserSpecifiedPrimaryKey returns true if a CREATE TABLE ... AS statement
+// has a PRIMARY KEY constraint specified.
+func (node *CreateTable) AsHasUserSpecifiedPrimaryKey() bool {
+	if node.As() {
+		for _, def := range node.Defs {
+			if d, ok := def.(*ColumnTableDef); !ok {
+				return false
+			} else if d.PrimaryKey {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // Format implements the NodeFormatter interface.
 func (node *CreateTable) Format(ctx *FmtCtx) {
 	ctx.WriteString("CREATE TABLE ")
@@ -914,9 +937,9 @@ func (node *CreateTable) Format(ctx *FmtCtx) {
 // but the CREATE TABLE tableName part.
 func (node *CreateTable) FormatBody(ctx *FmtCtx) {
 	if node.As() {
-		if len(node.AsColumnNames) > 0 {
+		if len(node.Defs) > 0 {
 			ctx.WriteString(" (")
-			ctx.FormatNode(&node.AsColumnNames)
+			ctx.FormatNode(&node.Defs)
 			ctx.WriteByte(')')
 		}
 		ctx.WriteString(" AS ")