Skip to content

Commit

Permalink
Merge #38589
Browse files Browse the repository at this point in the history 
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>
  • Loading branch information
@adityamaru27
craig[bot] and adityamaru27 committed Jul 16, 2019
2 parents 8809566 + f8333e7 commit 5dfa38c
Show file tree
Hide file tree
Showing 9 changed files with 385 additions and 84 deletions.
6 changes: 2 additions & 4 deletions docs/generated/sql/bnf/create_table_as_stmt.bnf
View file
Original file line number Diff line number Diff line change
@@ -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
32 changes: 30 additions & 2 deletions docs/generated/sql/bnf/stmt_block.bnf
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1352,6 +1352,10 @@ opt_table_elem_list ::=
table_elem_list
|

create_as_opt_col_list ::=
'(' create_as_table_defs ')'
|

view_name ::=
table_name

Expand Down Expand Up @@ -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 ')'

Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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 ')'
|
Expand Down Expand Up @@ -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
Expand Down
108 changes: 73 additions & 35 deletions pkg/sql/create_table.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
}

Expand All @@ -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 {
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
}

Expand Down
92 changes: 92 additions & 0 deletions pkg/sql/logictest/testdata/logic_test/create_as
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

30 changes: 19 additions & 11 deletions pkg/sql/opt/optbuilder/create_table.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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.
Expand Down
Loading

0 comments on commit 5dfa38c

Please sign in to comment.