diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go index 19dfa9868d09..fed7753ee714 100644 --- a/pkg/sql/alter_table.go +++ b/pkg/sql/alter_table.go @@ -141,10 +141,6 @@ func (n *alterTableNode) startExec(params runParams) error { switch t := cmd.(type) { case *tree.AlterTableAddColumn: d := t.ColumnDef - if d.HasFKConstraint() { - return unimplemented.NewWithIssue(32917, - "adding a REFERENCES constraint while also adding a column via ALTER not supported") - } version := params.ExecCfg().Settings.Version.ActiveVersionOrEmpty(params.ctx) if supported, err := isTypeSupportedInVersion(version, d.Type); err != nil { return err @@ -317,12 +313,8 @@ func (n *alterTableNode) startExec(params runParams) error { case *tree.ForeignKeyConstraintTableDef: for _, colName := range d.FromCols { - col, err := n.tableDesc.FindActiveColumnByName(string(colName)) + col, err := n.tableDesc.FindActiveOrNewColumnByName(colName) if err != nil { - if _, dropped, inactiveErr := n.tableDesc.FindColumnByName(colName); inactiveErr == nil && !dropped { - return unimplemented.NewWithIssue(32917, - "adding a REFERENCES constraint while the column is being added not supported") - } return err } diff --git a/pkg/sql/create_table.go b/pkg/sql/create_table.go index 5cc26bfa4751..8c8e60f44c88 100644 --- a/pkg/sql/create_table.go +++ b/pkg/sql/create_table.go @@ -608,16 +608,16 @@ func ResolveFK( validationBehavior tree.ValidationBehavior, settings *cluster.Settings, ) error { - originColumnIDs := make(sqlbase.ColumnIDs, len(d.FromCols)) + originCols := make([]*sqlbase.ColumnDescriptor, len(d.FromCols)) for i, col := range d.FromCols { - col, _, err := tbl.FindColumnByName(col) + col, err := tbl.FindActiveOrNewColumnByName(col) if err != nil { return err } if err := col.CheckCanBeFKRef(); err != nil { return err } - originColumnIDs[i] = col.ID + originCols[i] = col } target, err := ResolveMutableExistingObject(ctx, sc, &d.Table, true /*required*/, ResolveRequireTableDesc) @@ -658,33 +658,28 @@ func ResolveFK( } } - srcCols, err := tbl.FindActiveColumnsByNames(d.FromCols) - if err != nil { - return err - } - - targetColNames := d.ToCols + referencedColNames := d.ToCols // If no columns are specified, attempt to default to PK. - if len(targetColNames) == 0 { - targetColNames = make(tree.NameList, len(target.PrimaryIndex.ColumnNames)) + if len(referencedColNames) == 0 { + referencedColNames = make(tree.NameList, len(target.PrimaryIndex.ColumnNames)) for i, n := range target.PrimaryIndex.ColumnNames { - targetColNames[i] = tree.Name(n) + referencedColNames[i] = tree.Name(n) } } - targetCols, err := target.FindActiveColumnsByNames(targetColNames) + referencedCols, err := target.FindActiveColumnsByNames(referencedColNames) if err != nil { return err } - if len(targetCols) != len(srcCols) { + if len(referencedCols) != len(originCols) { return pgerror.Newf(pgcode.Syntax, "%d columns must reference exactly %d columns in referenced table (found %d)", - len(srcCols), len(srcCols), len(targetCols)) + len(originCols), len(originCols), len(referencedCols)) } - for i := range srcCols { - if s, t := srcCols[i], targetCols[i]; !s.Type.Equivalent(&t.Type) { + for i := range originCols { + if s, t := originCols[i], referencedCols[i]; !s.Type.Equivalent(&t.Type) { return pgerror.Newf(pgcode.DatatypeMismatch, "type of %q (%s) does not match foreign key %q.%q (%s)", s.Name, s.Type.String(), target.Name, t.Name, t.Type.String()) @@ -715,17 +710,17 @@ func ResolveFK( } } - targetColIDs := make(sqlbase.ColumnIDs, len(targetCols)) - for i := range targetCols { - targetColIDs[i] = targetCols[i].ID + targetColIDs := make(sqlbase.ColumnIDs, len(referencedCols)) + for i := range referencedCols { + targetColIDs[i] = referencedCols[i].ID } // Don't add a SET NULL action on an index that has any column that is NOT // NULL. if d.Actions.Delete == tree.SetNull || d.Actions.Update == tree.SetNull { - for _, sourceColumn := range srcCols { - if !sourceColumn.Nullable { - col := qualifyFKColErrorWithDB(ctx, txn, tbl.TableDesc(), sourceColumn.Name) + for _, originColumn := range originCols { + if !originColumn.Nullable { + col := qualifyFKColErrorWithDB(ctx, txn, tbl.TableDesc(), originColumn.Name) return pgerror.Newf(pgcode.InvalidForeignKey, "cannot add a SET NULL cascading action on column %q which has a NOT NULL constraint", col, ) @@ -736,11 +731,11 @@ func ResolveFK( // Don't add a SET DEFAULT action on an index that has any column that has // a DEFAULT expression of NULL and a NOT NULL constraint. if d.Actions.Delete == tree.SetDefault || d.Actions.Update == tree.SetDefault { - for _, sourceColumn := range srcCols { + for _, originColumn := range originCols { // Having a default expression of NULL, and a constraint of NOT NULL is a // contradiction and should never be allowed. - if sourceColumn.DefaultExpr == nil && !sourceColumn.Nullable { - col := qualifyFKColErrorWithDB(ctx, txn, tbl.TableDesc(), sourceColumn.Name) + if originColumn.DefaultExpr == nil && !originColumn.Nullable { + col := qualifyFKColErrorWithDB(ctx, txn, tbl.TableDesc(), originColumn.Name) return pgerror.Newf(pgcode.InvalidForeignKey, "cannot add a SET DEFAULT cascading action on column %q which has a "+ "NOT NULL constraint and a NULL default expression", col, @@ -749,10 +744,15 @@ func ResolveFK( } } + originColumnIDs := make(sqlbase.ColumnIDs, len(originCols)) + for i, col := range originCols { + originColumnIDs[i] = col.ID + } var legacyOriginIndexID sqlbase.IndexID // Search for an index on the origin table that matches. If one doesn't exist, - // we create one automatically if the table to alter is new or empty. - originIdx, err := sqlbase.FindFKOriginIndex(tbl.TableDesc(), originColumnIDs) + // we create one automatically if the table to alter is new or empty. We also + // search if an index for the set of columns was created in this transaction. + originIdx, err := sqlbase.FindFKOriginIndexInTxn(tbl, originColumnIDs) if err == nil { // If there was no error, we found a suitable index. legacyOriginIndexID = originIdx.ID @@ -775,7 +775,7 @@ func ResolveFK( return pgerror.Newf(pgcode.ForeignKeyViolation, "foreign key requires an existing index on columns %s", colNames.String()) } - id, err := addIndexForFK(tbl, srcCols, constraintName, ts) + id, err := addIndexForFK(tbl, originCols, constraintName, ts) if err != nil { return err } @@ -825,7 +825,7 @@ func ResolveFK( // that will support using `srcCols` as the referencing (src) side of an FK. func addIndexForFK( tbl *sqlbase.MutableTableDescriptor, - srcCols []sqlbase.ColumnDescriptor, + srcCols []*sqlbase.ColumnDescriptor, constraintName string, ts FKTableState, ) (sqlbase.IndexID, error) { diff --git a/pkg/sql/logictest/testdata/logic_test/alter_table b/pkg/sql/logictest/testdata/logic_test/alter_table index ea58e856812f..763e278fe108 100644 --- a/pkg/sql/logictest/testdata/logic_test/alter_table +++ b/pkg/sql/logictest/testdata/logic_test/alter_table @@ -381,9 +381,6 @@ ALTER TABLE t ADD i INT DEFAULT 1 CHECK (i < g) statement error pq: validation of CHECK "i > 0" failed on row:.* g=1.* i=0 ALTER TABLE t ADD i INT AS (g - 1) STORED CHECK (i > 0) -statement error adding a REFERENCES constraint while also adding a column via ALTER not supported -ALTER TABLE t ADD f INT UNIQUE REFERENCES other - query TTTTB SHOW CONSTRAINTS FROM t ---- @@ -926,18 +923,6 @@ CREATE TABLE t32917_2 (b INT PRIMARY KEY) statement ok INSERT INTO t32917_2 VALUES (1), (2), (3) -statement ok -BEGIN - -statement ok -ALTER TABLE t32917_2 ADD c INT UNIQUE DEFAULT 4 - -statement error adding a REFERENCES constraint while the column is being added not supported -ALTER TABLE t32917_2 ADD CONSTRAINT fk_c_a FOREIGN KEY (c) references t32917 (a) - -statement ok -ROLLBACK - # Test SET NOT NULL statement ok CREATE TABLE t (a INT) @@ -1135,3 +1120,250 @@ CREATE TABLE t25045 (x INT, y INT AS (x+1) STORED) statement error pq: column \"x\" is referenced by computed column \"y\" ALTER TABLE t25045 DROP COLUMN x + +subtest add_col_references + +statement ok +DROP TABLE IF EXISTS t1, t2; +CREATE TABLE t1 (x INT PRIMARY KEY); +CREATE TABLE t2 (y INT) + +statement ok +ALTER TABLE t2 ADD COLUMN x INT REFERENCES t1 (x) + +statement ok +INSERT INTO t1 VALUES (1) + +statement error pq: insert on table "t2" violates foreign key constraint "fk_x_ref_t1" +INSERT INTO t2 VALUES (2, 2) + +statement ok +INSERT INTO t2 VALUES (1, 1) + +# Error out trying to add a column with a foreign key on a non-empty table. +statement error pq: foreign key requires an existing index on columns \("z"\) +ALTER TABLE t2 ADD COLUMN z INT REFERENCES t1 (x) + +# Check that the foreign key was indeed added. +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + y INT8 NULL, + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + INDEX t2_auto_index_fk_x_ref_t1 (x ASC), + FAMILY "primary" (y, rowid, x) +) + +# Test that only one index gets created when adding a column +# with references and unique. +statement ok +CREATE TABLE t3 (y INT) + +statement ok +ALTER TABLE t3 ADD COLUMN x INT UNIQUE REFERENCES t1 (x) + +query TT +SHOW CREATE t3 +---- +t3 CREATE TABLE t3 ( + y INT8 NULL, + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + UNIQUE INDEX t3_x_key (x ASC), + FAMILY "primary" (y, rowid, x) +) + +# We allowed the foreign key validation code to look into the mutations +# list to validate what columns / indexes can be used for foreign keys. +# Ensure that we still have the correct restrictions. +statement ok +DROP TABLE t1, t2 CASCADE; +CREATE TABLE t1 (x INT PRIMARY KEY); +CREATE TABLE t2 (x INT, y INT, INDEX i (x)) + +statement error pq: column \"x\" does not exist +BEGIN; +ALTER TABLE t2 DROP COLUMN x; +ALTER TABLE t2 ADD FOREIGN KEY (x) REFERENCES t1(x); + +statement ok +ROLLBACK + +statement ok +INSERT INTO t2 VALUES (1, 2) + +statement error pq: foreign key requires an existing index on columns \("x"\) +BEGIN; +DROP INDEX t2@i; +ALTER TABLE t2 ADD FOREIGN KEY (x) REFERENCES t1(x) + +statement ok +ROLLBACK + +# Test using ADD COL REFERENCES in a self referencing constraint. +statement ok +DROP TABLE t1 CASCADE; +CREATE TABLE t1 (x INT PRIMARY KEY); +ALTER TABLE t1 ADD COLUMN x2 INT REFERENCES t1 (x) + +query TT +SHOW CREATE t1 +---- +t1 CREATE TABLE t1 ( + x INT8 NOT NULL, + x2 INT8 NULL, + CONSTRAINT "primary" PRIMARY KEY (x ASC), + CONSTRAINT fk_x2_ref_t1 FOREIGN KEY (x2) REFERENCES t1(x), + INDEX t1_auto_index_fk_x2_ref_t1 (x2 ASC), + FAMILY "primary" (x, x2) +) + +statement error pq: insert on table "t1" violates foreign key constraint "fk_x2_ref_t1" +INSERT INTO t1 VALUES (1, 2) + +# Test ADD COL REFERENCES on a new table in the same txn. +statement ok +DROP TABLE t1, t2 CASCADE + +statement ok +BEGIN; +CREATE TABLE t1 (x INT PRIMARY KEY); +CREATE TABLE t2 (y INT); +ALTER TABLE t2 ADD COLUMN x INT REFERENCES t1 (x); +COMMIT + +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + y INT8 NULL, + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + INDEX t2_auto_index_fk_x_ref_t1 (x ASC), + FAMILY "primary" (y, rowid, x) +) + +# Test that we can also add a column and then an FK in the same txn. +statement ok +DROP TABLE t1, t2 CASCADE + +statement ok +BEGIN; +CREATE TABLE t1 (x INT PRIMARY KEY); +CREATE TABLE t2 (y INT); +ALTER TABLE t2 ADD COLUMN x INT; +ALTER TABLE t2 ADD FOREIGN KEY (x) REFERENCES t1 (x); +COMMIT + +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + y INT8 NULL, + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + INDEX t2_auto_index_fk_x_ref_t1 (x ASC), + FAMILY "primary" (y, rowid, x) +) + +# Test that we can add a column and an index to an FK in the same txn. +statement ok +DROP TABLE t1, t2 CASCADE + +statement ok +BEGIN; +CREATE TABLE t1 (x INT PRIMARY KEY); +INSERT INTO t1 VALUES (1); +CREATE TABLE t2 (y INT); +INSERT INTO t2 VALUES (2); +ALTER TABLE t2 ADD COLUMN x INT; +CREATE INDEX ON t2 (x); +ALTER TABLE t2 ADD FOREIGN KEY (x) REFERENCES t1 (x); +COMMIT + +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + y INT8 NULL, + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + INDEX t2_x_idx (x ASC), + FAMILY "primary" (y, rowid, x) +) + +# Test the above on a table not created in the same txn. +statement ok +DROP TABLE t1, t2 CASCADE; +CREATE TABLE t1 (x INT PRIMARY KEY); +CREATE TABLE t2 (y INT) + +statement ok +BEGIN; +ALTER TABLE t2 ADD COLUMN x INT; +CREATE INDEX ON t2 (x); +ALTER TABLE t2 ADD FOREIGN KEY (x) REFERENCES t1 (x); +COMMIT + +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + y INT8 NULL, + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + INDEX t2_x_idx (x ASC), + FAMILY "primary" (y, rowid, x) +) + +# Test that an FK can use a newly created index. +statement ok +DROP TABLE t1, t2 CASCADE; +CREATE TABLE t1 (x INT PRIMARY KEY); +CREATE TABLE t2 (x INT) + +statement ok +BEGIN; +CREATE INDEX ON t2 (x); +ALTER TABLE t2 ADD FOREIGN KEY (x) REFERENCES t1 (x); +COMMIT + +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + x INT8 NULL, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + INDEX t2_x_idx (x ASC), + FAMILY "primary" (x, rowid) +) + +# Test when default column value leads to an FK violation. +statement ok +DROP TABLE t1, t2 CASCADE + +statement ok +CREATE TABLE t1 (x INT PRIMARY KEY); +INSERT INTO t1 VALUES (1); +CREATE TABLE t2 (y INT); +INSERT INTO t2 VALUES (2) + +statement error pq: foreign key violation +ALTER TABLE t2 ADD COLUMN x INT DEFAULT 2 UNIQUE REFERENCES t1 (x) + +# Test that it works with an appropriate default. +statement ok +ALTER TABLE t2 ADD COLUMN x INT DEFAULT 1 UNIQUE REFERENCES t1 (x) + +query TT +SHOW CREATE t2 +---- +t2 CREATE TABLE t2 ( + y INT8 NULL, + x INT8 NULL DEFAULT 1:::INT8, + CONSTRAINT fk_x_ref_t1 FOREIGN KEY (x) REFERENCES t1(x), + UNIQUE INDEX t2_x_key (x ASC), + FAMILY "primary" (y, rowid, x) +) diff --git a/pkg/sql/logictest/testdata/logic_test/fk b/pkg/sql/logictest/testdata/logic_test/fk index c52cdb7f565b..a2a09e085157 100644 --- a/pkg/sql/logictest/testdata/logic_test/fk +++ b/pkg/sql/logictest/testdata/logic_test/fk @@ -2525,12 +2525,6 @@ BEGIN; ALTER TABLE parentid ADD id INT NOT NULL AS (k + 2) STORED; ALTER TABLE c statement ok ROLLBACK; -statement error adding a REFERENCES constraint while the column is being added not supported -BEGIN; ALTER TABLE childid ADD id2 INT UNIQUE NOT NULL DEFAULT 0; ALTER TABLE childid ADD CONSTRAINT fk_id FOREIGN KEY (id2) REFERENCES parentid (k); - -statement ok -ROLLBACK; - subtest dont_check_nulls # Make sure that nulls are never checked while executing FK constraints. diff --git a/pkg/sql/logictest/testdata/logic_test/fk_opt b/pkg/sql/logictest/testdata/logic_test/fk_opt index d2288f4c2d9e..a94c914d5dd1 100644 --- a/pkg/sql/logictest/testdata/logic_test/fk_opt +++ b/pkg/sql/logictest/testdata/logic_test/fk_opt @@ -2769,12 +2769,6 @@ BEGIN; ALTER TABLE parentid ADD id INT NOT NULL AS (k + 2) STORED; ALTER TABLE c statement ok ROLLBACK; -statement error adding a REFERENCES constraint while the column is being added not supported -BEGIN; ALTER TABLE childid ADD id2 INT UNIQUE NOT NULL DEFAULT 0; ALTER TABLE childid ADD CONSTRAINT fk_id FOREIGN KEY (id2) REFERENCES parentid (k); - -statement ok -ROLLBACK; - subtest dont_check_nulls # Make sure that nulls are never checked while executing FK constraints. diff --git a/pkg/sql/sem/tree/alter_table.go b/pkg/sql/sem/tree/alter_table.go index 123a54e2150c..4cc73f226e8c 100644 --- a/pkg/sql/sem/tree/alter_table.go +++ b/pkg/sql/sem/tree/alter_table.go @@ -161,6 +161,26 @@ func (node *AlterTable) HoistAddColumnConstraints() { ) } d.CheckExprs = nil + if d.HasFKConstraint() { + var targetCol NameList + if d.References.Col != "" { + targetCol = append(targetCol, d.References.Col) + } + fk := &ForeignKeyConstraintTableDef{ + Table: *d.References.Table, + FromCols: NameList{d.Name}, + ToCols: targetCol, + Name: d.References.ConstraintName, + Actions: d.References.Actions, + Match: d.References.Match, + } + constraint := &AlterTableAddConstraint{ + ConstraintDef: fk, + ValidationBehavior: ValidationDefault, + } + normalizedCmds = append(normalizedCmds, constraint) + d.References.Table = nil + } } } node.Cmds = normalizedCmds diff --git a/pkg/sql/sqlbase/structured.go b/pkg/sql/sqlbase/structured.go index e792bd65e7bf..310247d57be1 100644 --- a/pkg/sql/sqlbase/structured.go +++ b/pkg/sql/sqlbase/structured.go @@ -2526,6 +2526,30 @@ func (desc *TableDescriptor) FindColumnByName(name tree.Name) (*ColumnDescriptor return nil, false, NewUndefinedColumnError(string(name)) } +// FindActiveOrNewColumnByName finds the column with the specified name. +// It returns either an active column or a column that was added in the +// same transaction that is currently running. +func (desc *MutableTableDescriptor) FindActiveOrNewColumnByName( + name tree.Name, +) (*ColumnDescriptor, error) { + for i := range desc.Columns { + c := &desc.Columns[i] + if c.Name == string(name) { + return c, nil + } + } + currentMutationID := desc.ClusterVersion.NextMutationID + for i := range desc.Mutations { + mut := &desc.Mutations[i] + if col := mut.GetColumn(); col != nil && + mut.MutationID == currentMutationID && + mut.Direction == DescriptorMutation_ADD { + return col, nil + } + } + return nil, NewUndefinedColumnError(string(name)) +} + // ColumnIdxMap returns a map from Column ID to the ordinal position of that // column. func (desc *TableDescriptor) ColumnIdxMap() map[ColumnID]int { diff --git a/pkg/sql/sqlbase/table.go b/pkg/sql/sqlbase/table.go index 0fd75f59d445..1986ec1e35f5 100644 --- a/pkg/sql/sqlbase/table.go +++ b/pkg/sql/sqlbase/table.go @@ -532,6 +532,43 @@ func FindFKOriginIndex( ) } +// FindFKOriginIndexInTxn finds the first index in the supplied originTable +// that can satisfy an outgoing foreign key of the supplied column ids. +// It returns either an index that is active, or an index that was created +// in the same transaction that is currently running. +func FindFKOriginIndexInTxn( + originTable *MutableTableDescriptor, originColIDs ColumnIDs, +) (*IndexDescriptor, error) { + // Search for an index on the origin table that matches our foreign + // key columns. + if originTable.PrimaryIndex.IsValidOriginIndex(originColIDs) { + return &originTable.PrimaryIndex, nil + } + // If the PK doesn't match, find the index corresponding to the origin column. + for i := range originTable.Indexes { + idx := &originTable.Indexes[i] + if idx.IsValidOriginIndex(originColIDs) { + return idx, nil + } + } + currentMutationID := originTable.ClusterVersion.NextMutationID + for i := range originTable.Mutations { + mut := &originTable.Mutations[i] + if idx := mut.GetIndex(); idx != nil && + mut.MutationID == currentMutationID && + mut.Direction == DescriptorMutation_ADD { + if idx.IsValidOriginIndex(originColIDs) { + return idx, nil + } + } + } + return nil, pgerror.Newf( + pgcode.ForeignKeyViolation, + "there is no index matching given keys for referenced table %s", + originTable.Name, + ) +} + // ConditionalGetTableDescFromTxn validates that the supplied TableDescriptor // matches the one currently stored in kv. This simulates a CPut and returns a // ConditionFailedError on mismatch. We don't directly use CPut with protos