Merge #38589

38589: pkg: Allow user defined primary key in CREATE TABLE ... AS r=adityamaru27 a=adityamaru27

This change allows users to specify which column in the newly created table should serve as the primary key (#20940). If no such constraint is specified, we generate a unique row_id to serve as the primary key.

Co-authored-by: Aditya Maru <adityamaru@cockroachlabs.com>

Author: craig[bot]

docs/generated/sql/bnf/create_table_as_stmt.bnf

@@ -1,5 +1,3 @@
 create_table_as_stmt ::=
-	'CREATE' 'TABLE' table_name '(' name ( ( ',' name ) )* ')' 'AS' select_stmt
-	| 'CREATE' 'TABLE' table_name  'AS' select_stmt
-	| 'CREATE' 'TABLE' 'IF' 'NOT' 'EXISTS' table_name '(' name ( ( ',' name ) )* ')' 'AS' select_stmt
-	| 'CREATE' 'TABLE' 'IF' 'NOT' 'EXISTS' table_name  'AS' select_stmt
+	'CREATE' 'TABLE' table_name create_as_opt_col_list 'AS' select_stmt
+	| 'CREATE' 'TABLE' 'IF' 'NOT' 'EXISTS' table_name create_as_opt_col_list 'AS' select_stmt

docs/generated/sql/bnf/stmt_block.bnf

@@ -969,8 +969,8 @@ create_table_stmt ::=
 	| 'CREATE' 'TABLE' 'IF' 'NOT' 'EXISTS' table_name '(' opt_table_elem_list ')' opt_interleave opt_partition_by
 
 create_table_as_stmt ::=
-	'CREATE' 'TABLE' table_name opt_column_list 'AS' select_stmt
-	| 'CREATE' 'TABLE' 'IF' 'NOT' 'EXISTS' table_name opt_column_list 'AS' select_stmt
+	'CREATE' 'TABLE' table_name create_as_opt_col_list 'AS' select_stmt
+	| 'CREATE' 'TABLE' 'IF' 'NOT' 'EXISTS' table_name create_as_opt_col_list 'AS' select_stmt
 
 create_view_stmt ::=
 	'CREATE' 'VIEW' view_name opt_column_list 'AS' select_stmt
@@ -1352,6 +1352,10 @@ opt_table_elem_list ::=
 	table_elem_list
 	| 
 
+create_as_opt_col_list ::=
+	'(' create_as_table_defs ')'
+	| 
+
 view_name ::=
 	table_name
 
@@ -1689,6 +1693,9 @@ partition_by ::=
 	| 'PARTITION' 'BY' 'RANGE' '(' name_list ')' '(' range_partitions ')'
 	| 'PARTITION' 'BY' 'NOTHING'
 
+create_as_table_defs ::=
+	( column_name create_as_col_qual_list ) ( ( ',' create_as_constraint_def | ',' column_name create_as_col_qual_list ) )*
+
 common_table_expr ::=
 	table_alias_name opt_column_list 'AS' '(' preparable_stmt ')'
 
@@ -1911,6 +1918,12 @@ list_partitions ::=
 range_partitions ::=
 	( range_partition ) ( ( ',' range_partition ) )*
 
+create_as_col_qual_list ::=
+	(  ) ( ( create_as_col_qualification ) )*
+
+create_as_constraint_def ::=
+	create_as_constraint_elem
+
 index_flags_param ::=
 	'FORCE_INDEX' '=' index_name
 	| 'NO_INDEX_JOIN'
@@ -2121,6 +2134,12 @@ list_partition ::=
 range_partition ::=
 	partition 'VALUES' 'FROM' '(' expr_list ')' 'TO' '(' expr_list ')' opt_partition_by
 
+create_as_col_qualification ::=
+	create_as_col_qualification_elem
+
+create_as_constraint_elem ::=
+	'PRIMARY' 'KEY' '(' create_as_params ')'
+
 col_qualification_elem ::=
 	'NOT' 'NULL'
 	| 'NULL'
@@ -2230,6 +2249,12 @@ func_expr_windowless ::=
 rowsfrom_list ::=
 	( rowsfrom_item ) ( ( ',' rowsfrom_item ) )*
 
+create_as_col_qualification_elem ::=
+	'PRIMARY' 'KEY'
+
+create_as_params ::=
+	( create_as_param ) ( ( ',' create_as_param ) )*
+
 opt_name_parens ::=
 	'(' name ')'
 	| 
@@ -2295,6 +2320,9 @@ join_outer ::=
 rowsfrom_item ::=
 	func_expr_windowless
 
+create_as_param ::=
+	column_name
+
 frame_extent ::=
 	frame_bound
 	| 'BETWEEN' frame_bound 'AND' frame_bound

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
 
@@ -985,38 +1010,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

@@ -42,7 +42,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(sqlbase.NewSyntaxError(fmt.Sprintf(
@@ -51,17 +56,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.