sql/internal/specuti: connect columns to their indexes on eval (#3223)

ariga · Nov 17, 2024 · 194c365 · 194c365
1 parent 147520c
commit 194c365
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 8 deletions.
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
@@ -477,16 +477,21 @@ func Index(spec *sqlspec.Index, parent *schema.Table, partFns ...func(*sqlspec.I
 			parts = append(parts, part)
 		}
 	}
-	i := &schema.Index{
+	idx := &schema.Index{
 		Name:   spec.Name,
 		Unique: spec.Unique,
 		Table:  parent,
 		Parts:  parts,
 	}
-	if err := convertCommentFromSpec(spec, &i.Attrs); err != nil {
+	if err := convertCommentFromSpec(spec, &idx.Attrs); err != nil {
 		return nil, err
 	}
-	return i, nil
+	for _, p := range idx.Parts {
+		if p.C != nil {
+			p.C.AddIndexes(idx)
+		}
+	}
+	return idx, nil
 }
 
 // Check converts a sqlspec.Check to a schema.Check.
@@ -510,10 +515,19 @@ func PrimaryKey(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index,
 			C:     c,
 		})
 	}
-	return &schema.Index{
+	pk := &schema.Index{
 		Table: parent,
 		Parts: parts,
-	}, nil
+	}
+	if err := convertCommentFromSpec(spec, &pk.Attrs); err != nil {
+		return nil, err
+	}
+	for _, p := range pk.Parts {
+		if p.C != nil {
+			p.C.AddIndexes(pk)
+		}
+	}
+	return pk, nil
 }
 
 // linkForeignKeys creates the foreign keys defined in the Table's spec by creating references

diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
@@ -281,8 +281,11 @@ schema "schema" {
 			{SeqNo: 0, C: exp.Tables[1].Columns[0]},
 		},
 	}
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].PrimaryKey)
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].Indexes[0])
+	exp.Tables[0].Columns[1].AddIndexes(exp.Tables[0].Indexes[0])
+	exp.Tables[1].Columns[0].AddIndexes(exp.Tables[1].PrimaryKey)
 	require.EqualValues(t, exp, &s)
-
 	_, err = MarshalHCL(&s)
 	require.NoError(t, err)
 }
@@ -708,7 +711,6 @@ table "users" {
 						),
 				),
 		)
-	exp.Tables[0].Columns[0].Indexes = nil
 	schema.NewRealm(exp)
 	require.EqualValues(t, exp, &s)
 }
@@ -801,7 +803,6 @@ schema "test" {
 						AddAttrs(&IndexType{T: IndexTypeFullText}, &IndexParser{P: "custom"}),
 				),
 		)
-	c.Indexes = nil
 	schema.NewRealm(exp)
 	require.EqualValues(t, exp, &s)
 }
@@ -870,6 +871,7 @@ schema "test" {
 		Parts: []*schema.IndexPart{{C: exp.Tables[0].Columns[0]}},
 		Attrs: []schema.Attr{&IndexType{T: IndexTypeHash}},
 	})
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].PrimaryKey)
 	schema.NewRealm(exp)
 	require.EqualValues(t, exp, &s)
 }

diff --git a/sql/postgres/sqlspec_oss_test.go b/sql/postgres/sqlspec_oss_test.go
@@ -277,6 +277,10 @@ enum "account_type" {
 			{SeqNo: 0, C: exp.Tables[1].Columns[0]},
 		},
 	}
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].PrimaryKey)
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].Indexes...)
+	exp.Tables[0].Columns[1].AddIndexes(exp.Tables[0].Indexes...)
+	exp.Tables[1].Columns[0].AddIndexes(exp.Tables[1].PrimaryKey)
 	exp.Realm = schema.NewRealm(exp)
 	require.EqualValues(t, exp, &s)
 }

diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
@@ -751,6 +751,16 @@ func (c *Column) SetPos(p *Pos) *Column {
 	return c
 }
 
+// AddIndexes appends the references to the indexes this column is part of.
+func (c *Column) AddIndexes(indexes ...*Index) *Column {
+	for _, idx := range indexes {
+		if !slices.Contains(c.Indexes, idx) {
+			c.Indexes = append(c.Indexes, idx)
+		}
+	}
+	return c
+}
+
 // NewCheck creates a new check.
 func NewCheck() *Check {
 	return &Check{}

diff --git a/sql/sqlite/sqlspec_test.go b/sql/sqlite/sqlspec_test.go
@@ -177,6 +177,10 @@ table "accounts" {
 			{SeqNo: 0, C: exp.Tables[1].Columns[0]},
 		},
 	}
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].PrimaryKey)
+	exp.Tables[0].Columns[0].AddIndexes(exp.Tables[0].Indexes[0])
+	exp.Tables[0].Columns[1].AddIndexes(exp.Tables[0].Indexes[0])
+	exp.Tables[1].Columns[0].AddIndexes(exp.Tables[1].PrimaryKey)
 	var s schema.Schema
 	err := EvalHCLBytes([]byte(f), &s, nil)
 	require.NoError(t, err)