Merge pull request #3283 from dolthub/elian/9316

dolthub/dolt#9316: Fix PK setting and add new fields for `CREATE TABLE ... SELECT`

Author: elianddb

enginetest/queries/create_table_queries.go

@@ -327,6 +327,148 @@ var CreateTableQueries = []WriteQueryTest{
 }
 
 var CreateTableScriptTests = []ScriptTest{
+	{
+		// https://github.com/dolthub/dolt/issues/9316
+		Name:         "CREATE TABLE with constraints AS SELECT osticket repro",
+		SkipPrepared: true, // SHOW KEYS with WHERE clause doesn't work with prepared statements
+		SetUpScript: []string{
+			"CREATE TABLE ost_form_entry (id INT PRIMARY KEY, object_id INT, object_type VARCHAR(1))",
+			"CREATE TABLE ost_form_entry_values (entry_id INT, field_id INT, value VARCHAR(100), value_id INT)",
+			"CREATE TABLE ost_form_field (id INT PRIMARY KEY)",
+			"INSERT INTO ost_form_entry VALUES (1, 100, 'U'), (2, 101, 'U'), (3, 102, 'X')",
+			"INSERT INTO ost_form_entry_values VALUES (1, 1, '[email protected]', 1000), (2, 1, '[email protected]', 1001), (3, 2, 'other', 2000)",
+			"INSERT INTO ost_form_field VALUES (1), (2)",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: `CREATE TABLE IF NOT EXISTS ost_user__cdata (
+						PRIMARY KEY (user_id)
+					) DEFAULT CHARSET=utf8 AS
+					SELECT
+						entry.object_id as user_id,
+						MAX(IF(field.id='1',coalesce(ans.value_id, ans.value),NULL)) as email
+						FROM ost_form_entry entry
+						JOIN ost_form_entry_values ans
+						ON ans.entry_id = entry.id
+						JOIN ost_form_field field
+						ON field.id=ans.field_id
+						WHERE entry.object_type='U' GROUP BY entry.object_id`,
+			},
+			{
+				Query: "SELECT * FROM ost_user__cdata ORDER BY user_id",
+				Expected: []sql.Row{
+					{100, "1000"},
+					{101, "1001"},
+				},
+			},
+			{
+				Query: "SHOW KEYS FROM ost_user__cdata WHERE Key_name = 'PRIMARY'",
+				Expected: []sql.Row{
+					{"ost_user__cdata", 0, "PRIMARY", 1, "user_id", nil, 0, nil, nil, "", "BTREE", "", "", "YES", nil},
+				},
+			},
+		},
+	},
+	{
+		// https://github.com/dolthub/dolt/issues/9316
+		Name:         "CREATE TABLE with constraints AS SELECT",
+		SkipPrepared: true,
+		SetUpScript: []string{
+			"CREATE TABLE t1 (a int not null, b varchar(10))",
+			"INSERT INTO t1 VALUES (1, 'one'), (2, 'two'), (3, 'three')",
+			"CREATE TABLE source (id int, name varchar(20))",
+			"INSERT INTO source VALUES (1, 'alice'), (2, 'bob'), (3, 'charlie')",
+			"CREATE TABLE override_src (a bigint, b int)",
+			"INSERT INTO override_src VALUES (100, 200)",
+			"CREATE TABLE base (a int, b varchar(10))",
+			"INSERT INTO base VALUES (1, 'alpha'), (2, 'beta')",
+			"CREATE TABLE multi_src (id int, email varchar(50), age int)",
+			"INSERT INTO multi_src VALUES (1, '[email protected]', 25), (2, '[email protected]', 30)",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "CREATE TABLE t2 (PRIMARY KEY(a)) SELECT * FROM t1",
+			},
+			{
+				Query: "SELECT * FROM t2 ORDER BY a",
+				Expected: []sql.Row{
+					{1, "one"},
+					{2, "two"},
+					{3, "three"},
+				},
+			},
+			{
+				Query: "SHOW KEYS FROM t2 WHERE Key_name = 'PRIMARY'",
+				Expected: []sql.Row{
+					{"t2", 0, "PRIMARY", 1, "a", nil, 0, nil, nil, "", "BTREE", "", "", "YES", nil},
+				},
+			},
+			{
+				Query: "CREATE TABLE indexed (KEY(name)) SELECT * FROM source",
+			},
+			{
+				Query: "SELECT * FROM indexed ORDER BY id",
+				Expected: []sql.Row{
+					{1, "alice"},
+					{2, "bob"},
+					{3, "charlie"},
+				},
+			},
+			{
+				Query: "SHOW KEYS FROM indexed WHERE Key_name = 'name'",
+				Expected: []sql.Row{
+					{"indexed", 1, "name", 1, "name", nil, 0, nil, nil, "YES", "BTREE", "", "", "YES", nil},
+				},
+			},
+			{
+				Query: "CREATE TABLE override (a TINYINT NOT NULL) SELECT a, b FROM override_src",
+			},
+			{
+				Query: "SELECT * FROM override",
+				Expected: []sql.Row{
+					{int8(100), 200},
+				},
+			},
+			{
+				Query: "SHOW CREATE TABLE override",
+				Expected: []sql.Row{
+					{"override", "CREATE TABLE `override` (\n  `a` tinyint NOT NULL,\n  `b` int\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"},
+				},
+			},
+			{
+				Query: "CREATE TABLE uniq (UNIQUE KEY(a)) SELECT * FROM base",
+			},
+			{
+				Query: "SELECT * FROM uniq ORDER BY a",
+				Expected: []sql.Row{
+					{1, "alpha"},
+					{2, "beta"},
+				},
+			},
+			{
+				Query: "SHOW KEYS FROM uniq WHERE Key_name = 'a'",
+				Expected: []sql.Row{
+					{"uniq", 0, "a", 1, "a", nil, 0, nil, nil, "YES", "BTREE", "", "", "YES", nil},
+				},
+			},
+			{
+				Query: "CREATE TABLE multi_idx (PRIMARY KEY(id), KEY(email), KEY(age)) SELECT * FROM multi_src",
+			},
+			{
+				Query: "SELECT * FROM multi_idx ORDER BY id",
+				Expected: []sql.Row{
+					{1, "[email protected]", 25},
+					{2, "[email protected]", 30},
+				},
+			},
+			{
+				Query: "SELECT COUNT(*) FROM information_schema.statistics WHERE table_name = 'multi_idx'",
+				Expected: []sql.Row{
+					{3},
+				},
+			},
+		},
+	},
 	{
 		// https://github.com/dolthub/dolt/issues/6682
 		Name: "display width for numeric types",

sql/analyzer/resolve_create_select.go

@@ -38,15 +38,35 @@ func resolveCreateSelect(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.
 		newSch[i] = &tempCol
 	}
 
-	pkOrdinals := make([]int, 0)
+	colNameToIdx := make(map[string]int, len(newSch))
 	for i, col := range newSch {
-		if col.PrimaryKey {
-			pkOrdinals = append(pkOrdinals, i)
+		colNameToIdx[col.Name] = i
+	}
+
+	// Apply primary key constraints from index definitions to the merged schema
+	var nonPkIndexes sql.IndexDefs
+	pkOrdinals := make([]int, 0)
+	for _, idx := range ct.Indexes() {
+		if idx.IsPrimary() {
+			for _, idxCol := range idx.Columns {
+				if i, ok := colNameToIdx[idxCol.Name]; ok {
+					// https://dev.mysql.com/doc/refman/8.0/en/create-table.html
+					newSch[i].PrimaryKey = true
+					newSch[i].Nullable = false
+					pkOrdinals = append(pkOrdinals, i)
+				}
+			}
+		} else {
+			nonPkIndexes = append(nonPkIndexes, idx)
 		}
 	}
 
 	newSpec := &plan.TableSpec{
-		Schema: sql.NewPrimaryKeySchema(newSch, pkOrdinals...),
+		Schema:    sql.NewPrimaryKeySchema(newSch, pkOrdinals...),
+		IdxDefs:   nonPkIndexes, // Only pass non-PK indexes since PK is in schema
+		FkDefs:    ct.ForeignKeys(),
+		ChDefs:    ct.Checks(),
+		Collation: ct.Collation,
 	}
 
 	newCreateTable := plan.NewCreateTable(ct.Database(), ct.Name(), ct.IfNotExists(), ct.Temporary(), newSpec)

sql/analyzer/validate_create_table.go

@@ -38,6 +38,13 @@ func validateCreateTable(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.
 		return nil, transform.SameTree, err
 	}
 
+	// For CREATE TABLE AS SELECT, skip validation here because the schema isn't complete yet.
+	// The resolveCreateSelect analyzer rule will merge schemas and create a new CreateTable node,
+	// which will be validated separately after the merge.
+	if ct.Select() != nil {
+		return n, transform.SameTree, nil
+	}
+
 	sch := ct.PkSchema().Schema
 	idxs := ct.Indexes()