ddl: support adding constraint with column in one spec (#36020)

ref #14766
pingcap · Jul 8, 2022 · 959073c · 959073c
1 parent 836f7a3
commit 959073c
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 5 deletions.
diff --git a/ddl/ddl_api.go b/ddl/ddl_api.go
@@ -2997,11 +2997,21 @@ func needToOverwriteColCharset(options []*ast.TableOption) bool {
 // `ALTER TABLE ADD COLUMN (c1 INT, c2 INT)` is split into
 // `ALTER TABLE ADD COLUMN c1 INT, ADD COLUMN c2 INT`.
 func resolveAlterTableAddColumns(spec *ast.AlterTableSpec) []*ast.AlterTableSpec {
-	specs := make([]*ast.AlterTableSpec, len(spec.NewColumns))
-	for i, col := range spec.NewColumns {
+	specs := make([]*ast.AlterTableSpec, 0, len(spec.NewColumns)+len(spec.NewConstraints))
+	for _, col := range spec.NewColumns {
 		t := *spec
 		t.NewColumns = []*ast.ColumnDef{col}
-		specs[i] = &t
+		t.NewConstraints = []*ast.Constraint{}
+		specs = append(specs, &t)
+	}
+	// Split the add constraints from AlterTableSpec.
+	for _, con := range spec.NewConstraints {
+		t := *spec
+		t.NewColumns = []*ast.ColumnDef{}
+		t.NewConstraints = []*ast.Constraint{}
+		t.Constraint = con
+		t.Tp = ast.AlterTableAddConstraint
+		specs = append(specs, &t)
 	}
 	return specs
 }
@@ -3019,7 +3029,7 @@ func resolveAlterTableSpec(ctx sessionctx.Context, specs []*ast.AlterTableSpec)
 		if isIgnorableSpec(spec.Tp) {
 			continue
 		}
-		if spec.Tp == ast.AlterTableAddColumns && len(spec.NewColumns) > 1 {
+		if spec.Tp == ast.AlterTableAddColumns && (len(spec.NewColumns) > 1 || len(spec.NewConstraints) > 0) {
 			validSpecs = append(validSpecs, resolveAlterTableAddColumns(spec)...)
 		} else {
 			validSpecs = append(validSpecs, spec)

diff --git a/ddl/multi_schema_change_test.go b/ddl/multi_schema_change_test.go
@@ -42,7 +42,7 @@ func TestMultiSchemaChangeAddColumns(t *testing.T) {
 	tk.MustExec("alter table t add column b int default 2, add column c int default 3;")
 	tk.MustQuery("select * from t;").Check(testkit.Rows("1 2 3"))
 
-	// Test add multiple columns in one spec.
+	// Test add multiple columns and constraints in one spec.
 	tk.MustExec("drop table if exists t;")
 	tk.MustExec("create table t (a int);")
 	tk.MustExec("insert into t values (1);")
@@ -55,6 +55,31 @@ func TestMultiSchemaChangeAddColumns(t *testing.T) {
 	tk.MustExec("alter table t add column (d int default 4, e int default 5);")
 	tk.MustQuery("select * from t;").Check(testkit.Rows("1 2 3 4 5"))
 
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (a int);")
+	tk.MustExec("insert into t values (1);")
+	tk.MustExec("alter table t add column (index i(a), index i1(a));")
+	tk.MustQuery("select * from t use index (i, i1);").Check(testkit.Rows("1"))
+
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (a int);")
+	tk.MustExec("insert into t values (1);")
+	tk.MustExec("alter table t add column (b int default 2, index i(a), primary key (a));")
+	tk.MustQuery("select * from t use index (i, primary)").Check(testkit.Rows("1 2"))
+
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (a int);")
+	tk.MustExec("insert into t values (1);")
+	tk.MustExec("alter table t add column (index i(a));")
+	tk.MustQuery("select * from t use index (i)").Check(testkit.Rows("1"))
+
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (a int, b int, c int);")
+	tk.MustExec("insert into t values (1, 2, 3);")
+	tk.MustExec("alter table t add column (index i1(a, b, c), index i2(c, b, a), index i3((a + 1)), index i4((c - 1)));")
+	tk.MustQuery("select * from t use index (i1, i2);").Check(testkit.Rows("1 2 3"))
+	tk.MustExec("admin check table t;")
+
 	tk.MustExec("drop table if exists t;")
 	tk.MustExec("create table t (a int default 1);")
 	tk.MustExec("insert into t values ();")