diff --git a/server/ast/alter_table.go b/server/ast/alter_table.go index afca2212da..f52706c4c6 100644 --- a/server/ast/alter_table.go +++ b/server/ast/alter_table.go @@ -98,6 +98,22 @@ func nodeAlterTableCmds( return nil, nil, err } vitessDdlCmds = append(vitessDdlCmds, statement) + + // Postgres (unlike MySQL) allows an inline FK constraint + // when altering a table to add a column. GMS doesn't support + // this directly in the ALTER TABLE ADD COLUMN DDL command, + // so we break this out into a separate DDL command. + if len(statement.TableSpec.Constraints) > 0 { + vitessDdlCmds = append(vitessDdlCmds, + &vitess.DDL{ + Action: "alter", + ConstraintAction: "add", + Table: tableName, + IfExists: ifExists, + TableSpec: statement.TableSpec, + }) + } + case *tree.AlterTableDropColumn: statement, err = nodeAlterTableDropColumn(ctx, cmd, tableName, ifExists) if err != nil { @@ -217,6 +233,14 @@ func nodeAlterTableAddColumn(ctx *Context, node *tree.AlterTableAddColumn, table tableSpec := &vitess.TableSpec{} tableSpec.AddColumn(vitessColumnDef) + if node.ColumnDef.References.Table != nil { + constraintDef, err := nodeForeignKeyDefinitionFromColumnTableDef(ctx, node.ColumnDef.Name, node.ColumnDef) + if err != nil { + return nil, err + } + tableSpec.AddConstraint(&vitess.ConstraintDefinition{Details: constraintDef}) + } + return &vitess.DDL{ Action: "alter", ColumnAction: "add", diff --git a/testing/go/alter_table_test.go b/testing/go/alter_table_test.go index 9bcab8d3ec..3fc653a375 100644 --- a/testing/go/alter_table_test.go +++ b/testing/go/alter_table_test.go @@ -491,6 +491,27 @@ func TestAlterTable(t *testing.T) { }, }, }, + { + Name: "ALTER TABLE ADD COLUMN with inline FK constraint", + SetUpScript: []string{ + "create table t (v varchar(100));", + "create table parent (id int primary key);", + }, + Assertions: []ScriptTestAssertion{ + { + Query: "ALTER TABLE t ADD COLUMN c1 int REFERENCES parent(id);", + Expected: []sql.Row{}, + }, + { + Query: "SELECT conname AS constraint_name, pg_get_constraintdef(oid) AS constraint_definition FROM pg_constraint WHERE conrelid = 't'::regclass AND contype='f';", + Expected: []sql.Row{{"t_c1_fkey", "FOREIGN KEY t_c1_fkey (c1) REFERENCES parent (id)"}}, + }, + { + Query: "INSERT INTO t VALUES ('abc', 123);", + ExpectedErr: "Foreign key violation on fk: `t_c1_fkey`", + }, + }, + }, { Name: "Rename table", SetUpScript: []string{