Repository: ariga/atlas
Pull request #3509 (2025-05-19 07:26:06)

sql/schema: add PosSetter to objects
Approved internally. 


Patch:
diff --git a/sql/internal/sqlx/dev.go b/sql/internal/sqlx/dev.go
index 905b5260..ffac7026 100644
--- a/sql/internal/sqlx/dev.go
+++ b/sql/internal/sqlx/dev.go
@@ -47,6 +47,7 @@ func (d *DevDriver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *sc
 			Schemas: make([]string, 0, len(r.Schemas)),
 		}
 	)
+	name2pos := make(key2pos)
 	for _, o := range r.Objects {
 		changes = append(changes, &schema.AddObject{
 			O: o,
@@ -54,8 +55,8 @@ func (d *DevDriver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *sc
 				&schema.IfNotExists{},
 			},
 		})
+		name2pos.putObject(o)
 	}
-	name2pos := make(key2pos)
 	for _, s := range r.Schemas {
 		k, _ := name2pos.put(s.Attrs, keyS, s.Name)
 		opts.Schemas = append(opts.Schemas, s.Name)
@@ -82,6 +83,7 @@ func (d *DevDriver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *sc
 			}
 		}
 		for _, o := range s.Objects {
+			name2pos.putObject(o, k)
 			changes = append(changes, &schema.AddObject{O: o})
 		}
 		for _, f := range s.Funcs {
@@ -166,12 +168,15 @@ func (d *DevDriver) NormalizeSchema(ctx context.Context, s *schema.Schema) (*sch
 		if d.PatchObject != nil {
 			d.PatchObject(s, o)
 		}
+		name2pos.putObject(o, k)
 		changes = append(changes, &schema.AddObject{O: o})
 	}
 	for _, f := range s.Funcs {
+		name2pos.put(f.Attrs, k, keyFn, f.Name)
 		changes = append(changes, &schema.AddFunc{F: f})
 	}
 	for _, p := range s.Procs {
+		name2pos.put(p.Attrs, k, keyPr, p.Name)
 		changes = append(changes, &schema.AddProc{P: p})
 	}
 	if err := d.Driver.ApplyChanges(ctx, changes, func(opts *migrate.PlanOptions) {
@@ -254,62 +259,97 @@ func (k key2pos) put(attrs []schema.Attr, typename ...string) (string, bool) {
 	return n, false
 }
 
+func (k key2pos) putObject(o schema.Object, prefix ...string) {
+	ns, ok := o.(schema.SpecTypeNamer)
+	if !ok {
+		return
+	}
+	t, n := ns.SpecType(), ns.SpecName()
+	if t == "" || n == "" {
+		return
+	}
+	nk := poskey(append(prefix, t, n)...)
+	if ps, ok := o.(schema.PosSetter); ok && ps.Pos() != nil {
+		k[nk] = ps.Pos()
+	}
+}
+
 func (k key2pos) patchRealm(r *schema.Realm) {
+	k.patchObjects(r.Objects)
 	for _, s := range r.Schemas {
 		k.patchSchema(s)
 	}
 }
 
+func (k key2pos) patchObjects(objs []schema.Object, prefix ...string) {
+	for _, o := range objs {
+		ns, ok := o.(schema.SpecTypeNamer)
+		if !ok {
+			continue
+		}
+		t, n := ns.SpecType(), ns.SpecName()
+		if t == "" || n == "" {
+			continue
+		}
+		st, ok := o.(schema.PosSetter)
+		if !ok {
+			continue
+		}
+		k.patch(st, append(prefix, t, n)...)
+	}
+}
+
 func (k key2pos) patchSchema(s *schema.Schema) {
-	ks, _ := k.patch(&s.Attrs, keyS, s.Name)
+	ks, _ := k.patch(s, keyS, s.Name)
+	k.patchObjects(s.Objects, ks)
 	for _, t := range s.Tables {
-		tk, ok := k.patch(&t.Attrs, ks, keyT, t.Name)
+		tk, ok := k.patch(t, ks, keyT, t.Name)
 		if !ok {
 			continue
 		}
 		for _, tr := range t.Triggers {
-			k.patch(&tr.Attrs, tk, keyTg, tr.Name)
+			k.patch(tr, tk, keyTg, tr.Name)
 		}
 		for _, c := range t.Columns {
-			k.patch(&c.Attrs, tk, keyC, c.Name)
+			k.patch(c, tk, keyC, c.Name)
 		}
 		for _, i := range t.Indexes {
-			k.patch(&i.Attrs, tk, keyI, i.Name)
+			k.patch(i, tk, keyI, i.Name)
 		}
 		for _, f := range t.ForeignKeys {
-			k.patch(&f.Attrs, tk, keyF, f.Symbol)
+			k.patch(f, tk, keyF, f.Symbol)
 		}
 		for _, ck := range t.Checks() {
-			k.patch(&ck.Attrs, tk, keyK, ck.Name)
+			k.patch(ck, tk, keyK, ck.Name)
 		}
 	}
 	for _, v := range s.Views {
-		vk, ok := k.patch(&v.Attrs, ks, keyV, v.Name)
+		vk, ok := k.patch(v, ks, keyV, v.Name)
 		if !ok {
 			continue
 		}
 		for _, tr := range v.Triggers {
-			k.patch(&tr.Attrs, vk, keyTg, tr.Name)
+			k.patch(tr, vk, keyTg, tr.Name)
 		}
 		for _, c := range v.Columns {
-			k.patch(&c.Attrs, vk, keyC, c.Name)
+			k.patch(c, vk, keyC, c.Name)
 		}
 		for _, i := range v.Indexes {
-			k.patch(&i.Attrs, vk, keyI, i.Name)
+			k.patch(i, vk, keyI, i.Name)
 		}
 	}
 	for _, f := range s.Funcs {
-		k.patch(&f.Attrs, ks, keyFn, f.Name)
+		k.patch(f, ks, keyFn, f.Name)
 	}
 	for _, p := range s.Procs {
-		k.patch(&p.Attrs, ks, keyPr, p.Name)
+		k.patch(p, ks, keyPr, p.Name)
 	}
 }
 
-func (k key2pos) patch(attrs *[]schema.Attr, typename ...string) (string, bool) {
+func (k key2pos) patch(ps schema.PosSetter, typename ...string) (string, bool) {
 	n := poskey(typename...)
 	if p, ok := k[n]; ok {
-		*attrs = append(*attrs, p)
+		ps.SetPos(p)
 		return n, true
 	}
 	return n, false
diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
index 39d9d847..7b0ff356 100644
--- a/sql/schema/dsl.go
+++ b/sql/schema/dsl.go
@@ -57,12 +57,6 @@ func (s *Schema) AddAttrs(attrs ...Attr) *Schema {
 	return s
 }
 
-// SetPos sets the position of the schema.
-func (s *Schema) SetPos(p *Pos) *Schema {
-	ReplaceOrAppend(&s.Attrs, p)
-	return s
-}
-
 // SetRealm sets the database/realm of the schema.
 func (s *Schema) SetRealm(r *Realm) *Schema {
 	s.Realm = r
@@ -258,12 +252,6 @@ func (t *Table) AddAttrs(attrs ...Attr) *Table {
 	return t
 }
 
-// SetPos sets the position of the table.
-func (t *Table) SetPos(p *Pos) *Table {
-	ReplaceOrAppend(&t.Attrs, p)
-	return t
-}
-
 // AddDeps adds the given objects as dependencies to the view.
 func (t *Table) AddDeps(objs ...Object) *Table {
 	t.Deps = append(t.Deps, objs...)
@@ -372,12 +360,6 @@ func (v *View) AddAttrs(attrs ...Attr) *View {
 	return v
 }
 
-// SetPos sets the position of the view.
-func (v *View) SetPos(p *Pos) *View {
-	ReplaceOrAppend(&v.Attrs, p)
-	return v
-}
-
 // AddDeps adds the given objects as dependencies to the view.
 func (v *View) AddDeps(objs ...Object) *View {
 	v.Deps = append(v.Deps, objs...)
@@ -745,12 +727,6 @@ func (c *Column) AddAttrs(attrs ...Attr) *Column {
 	return c
 }
 
-// SetPos sets the position of the column.
-func (c *Column) SetPos(p *Pos) *Column {
-	ReplaceOrAppend(&c.Attrs, p)
-	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 {
@@ -784,12 +760,6 @@ func (c *Check) AddAttrs(attrs ...Attr) *Check {
 	return c
 }
 
-// SetPos sets the position of the check.
-func (c *Check) SetPos(p *Pos) *Check {
-	ReplaceOrAppend(&c.Attrs, p)
-	return c
-}
-
 // NewIndex creates a new index with the given name.
 func NewIndex(name string) *Index {
 	return &Index{Name: name}
@@ -837,12 +807,6 @@ func (i *Index) AddAttrs(attrs ...Attr) *Index {
 	return i
 }
 
-// SetPos sets the position of the index.
-func (i *Index) SetPos(p *Pos) *Index {
-	ReplaceOrAppend(&i.Attrs, p)
-	return i
-}
-
 // AddColumns adds the columns to index parts.
 func (i *Index) AddColumns(columns ...*Column) *Index {
 	for _, c := range columns {
@@ -904,12 +868,6 @@ func (p *IndexPart) AddAttrs(attrs ...Attr) *IndexPart {
 	return p
 }
 
-// SetPos sets the position of the index part.
-func (p *IndexPart) SetPos(p1 *Pos) *IndexPart {
-	ReplaceOrAppend(&p.Attrs, p1)
-	return p
-}
-
 // SetColumn sets the column of the index-part.
 func (p *IndexPart) SetColumn(c *Column) *IndexPart {
 	p.C = c
@@ -984,12 +942,6 @@ func (f *ForeignKey) AddAttrs(attrs ...Attr) *ForeignKey {
 	return f
 }
 
-// SetPos sets the position of the foreign key.
-func (f *ForeignKey) SetPos(p *Pos) *ForeignKey {
-	ReplaceOrAppend(&f.Attrs, p)
-	return f
-}
-
 // AddDeps adds the given objects as dependencies to the function.
 func (f *Func) AddDeps(objs ...Object) *Func {
 	f.Deps = append(f.Deps, objs...)
@@ -1070,7 +1022,7 @@ func (p *Pos) SetEnd(e struct{ Line, Column, Byte int }) *Pos {
 }
 
 // del searches an attribute of the same type as v in
-// the list and delete it.
+// the list and deletes it.
 func del(attrs *[]Attr, v Attr) {
 	t := reflect.TypeOf(v)
 	for i := range *attrs {
diff --git a/sql/schema/schema.go b/sql/schema/schema.go
index f77d1dee..08ebc1d4 100644
--- a/sql/schema/schema.go
+++ b/sql/schema/schema.go
@@ -5,7 +5,9 @@
 package schema
 
 import (
+	"fmt"
 	"reflect"
+	"strings"
 )
 
 type (
@@ -243,14 +245,31 @@ func (r *Realm) Object(f func(Object) bool) (Object, bool) {
 	return nil, false
 }
 
+// PosSetter wraps the two methods for getting
+// and setting positions for schema objects.
+type PosSetter interface {
+	Pos() *Pos
+	SetPos(*Pos)
+}
+
 // Pos of the schema, if exists.
-func (s *Schema) Pos() (*Pos, bool) {
+func (s *Schema) Pos() *Pos {
 	for _, a := range s.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
+}
+
+// Pos of the enum, if exists.
+func (e *EnumType) Pos() *Pos {
+	for _, a := range e.Attrs {
+		if p, ok := a.(*Pos); ok {
+			return p
+		}
+	}
+	return nil
 }
 
 // Table returns the first table that matched the given name.
@@ -324,13 +343,13 @@ func (t *Table) Column(name string) (*Column, bool) {
 }
 
 // Pos of the table, if exists.
-func (t *Table) Pos() (*Pos, bool) {
+func (t *Table) Pos() *Pos {
 	for _, a := range t.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Index returns the first index that matched the given name.
@@ -374,13 +393,13 @@ func (t *Table) Checks() (ck []*Check) {
 }
 
 // Pos of the view, if exists.
-func (v *View) Pos() (*Pos, bool) {
+func (v *View) Pos() *Pos {
 	for _, a := range v.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Materialized reports if the view is materialized.
@@ -440,54 +459,129 @@ func (v *View) AsTable() *Table {
 		AddColumns(v.Columns...)
 }
 
+// SetPos sets the position of the function.
+func (f *Func) SetPos(p *Pos) {
+	ReplaceOrAppend(&f.Attrs, p)
+}
+
+// Pos of the schema, if exists.
+func (f *Func) Pos() *Pos {
+	for _, a := range f.Attrs {
+		if p, ok := a.(*Pos); ok {
+			return p
+		}
+	}
+	return nil
+}
+
+// SetPos sets the position of the schema.
+func (s *Schema) SetPos(p *Pos) {
+	ReplaceOrAppend(&s.Attrs, p)
+}
+
+// SetPos sets the position of the enum type.
+func (e *EnumType) SetPos(p *Pos) {
+	ReplaceOrAppend(&e.Attrs, p)
+}
+
+// SetPos sets the position of the table.
+func (t *Table) SetPos(p *Pos) {
+	ReplaceOrAppend(&t.Attrs, p)
+}
+
+// SetPos sets the position of the view.
+func (v *View) SetPos(p *Pos) {
+	ReplaceOrAppend(&v.Attrs, p)
+}
+
+// SetPos sets the position of the column.
+func (c *Column) SetPos(p *Pos) {
+	ReplaceOrAppend(&c.Attrs, p)
+}
+
+// SetPos sets the position of the check.
+func (c *Check) SetPos(p *Pos) {
+	ReplaceOrAppend(&c.Attrs, p)
+}
+
+// SetPos sets the position of the index.
+func (i *Index) SetPos(p *Pos) {
+	ReplaceOrAppend(&i.Attrs, p)
+}
+
+// SetPos sets the position of the index part.
+func (p *IndexPart) SetPos(p1 *Pos) {
+	ReplaceOrAppend(&p.Attrs, p1)
+}
+
+// SetPos sets the position of the foreign key.
+func (f *ForeignKey) SetPos(p *Pos) {
+	ReplaceOrAppend(&f.Attrs, p)
+}
+
+// SetPos sets the position of the procedure.
+func (p *Proc) SetPos(p1 *Pos) {
+	ReplaceOrAppend(&p.Attrs, p1)
+}
+
+// Pos of the schema, if exists.
+func (p *Proc) Pos() *Pos {
+	for _, a := range p.Attrs {
+		if p, ok := a.(*Pos); ok {
+			return p
+		}
+	}
+	return nil
+}
+
 // Pos of the column, if exists.
-func (c *Column) Pos() (*Pos, bool) {
+func (c *Column) Pos() *Pos {
 	for _, a := range c.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Pos of the index, if exists.
-func (i *Index) Pos() (*Pos, bool) {
+func (i *Index) Pos() *Pos {
 	for _, a := range i.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Pos of the index part, if exists.
-func (p *IndexPart) Pos() (*Pos, bool) {
+func (p *IndexPart) Pos() *Pos {
 	for _, a := range p.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Pos of the check, if exists.
-func (c *Check) Pos() (*Pos, bool) {
+func (c *Check) Pos() *Pos {
 	for _, a := range c.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Pos of the foreign-key, if exists.
-func (f *ForeignKey) Pos() (*Pos, bool) {
+func (f *ForeignKey) Pos() *Pos {
 	for _, a := range f.Attrs {
 		if p, ok := a.(*Pos); ok {
-			return p, true
+			return p
 		}
 	}
-	return nil, false
+	return nil
 }
 
 // Column returns the first column that matches the given name.
@@ -510,6 +604,21 @@ func (f *ForeignKey) RefColumn(name string) (*Column, bool) {
 	return nil, false
 }
 
+// SetPos sets the position of the trigger.
+func (t *Trigger) SetPos(p *Pos) {
+	ReplaceOrAppend(&t.Attrs, p)
+}
+
+// Pos of the schema, if exists.
+func (t *Trigger) Pos() *Pos {
+	for _, a := range t.Attrs {
+		if p, ok := a.(*Pos); ok {
+			return p
+		}
+	}
+	return nil
+}
+
 // ReferenceOption for constraint actions.
 type ReferenceOption string
 
@@ -545,9 +654,10 @@ type (
 		T      string   // Optional type.
 		Values []string // Enum values.
 		Schema *Schema  // Optional schema.
+		Attrs  []Attr   // Extra attributes.
 	}
 
-	// BinaryType represents a type that stores a binary data.
+	// BinaryType represents a type that stores binary data.
 	BinaryType struct {
 		T    string
 		Size *int
@@ -726,6 +836,17 @@ type (
 	}
 )
 
+// String returns the position in editor/LSP style.
+// Format: "filename:line[:c][-end_line[:end_c]]"
+func (p *Pos) String() string {
+	var b strings.Builder
+	fmt.Fprintf(&b, "%s:%d", p.Filename, p.Start.Line)
+	if p.Start.Column > 0 {
+		fmt.Fprintf(&b, ":%d", p.Start.Column)
+	}
+	return b.String()
+}
+
 // A list of known view check options.
 const (
 	ViewCheckOptionNone     = "NONE"


Test patch:
diff --git a/sql/internal/sqlx/dev_test.go b/sql/internal/sqlx/dev_test.go
index e8a1b901..ea417999 100644
--- a/sql/internal/sqlx/dev_test.go
+++ b/sql/internal/sqlx/dev_test.go
@@ -46,13 +46,13 @@ func TestDriver_NormalizeRealm(t *testing.T) {
 		).
 		AddTables(
 			schema.NewTable("t1").
-				SetPos(
+				AddAttrs(
 					schema.NewFilePos("schema.hcl").
 						SetStart(hcl.Pos{Line: 2, Column: 2, Byte: 2}),
 				).
 				AddColumns(
 					schema.NewIntColumn("id", "int").
-						SetPos(
+						AddAttrs(
 							schema.NewFilePos("schema.hcl").
 								SetStart(hcl.Pos{Line: 3, Column: 3, Byte: 3}),
 						),
@@ -63,14 +63,11 @@ func TestDriver_NormalizeRealm(t *testing.T) {
 	)
 	normal, err = dev.NormalizeRealm(context.Background(), r)
 	require.NoError(t, err)
-	p, ok := normal.Schemas[0].Pos()
-	require.True(t, ok)
+	p := normal.Schemas[0].Pos()
 	require.Equal(t, schema.NewFilePos("schema.hcl").SetStart(hcl.Pos{Line: 1, Column: 1, Byte: 1}), p)
-	p, ok = normal.Schemas[0].Tables[0].Pos()
-	require.True(t, ok)
+	p = normal.Schemas[0].Tables[0].Pos()
 	require.Equal(t, schema.NewFilePos("schema.hcl").SetStart(hcl.Pos{Line: 2, Column: 2, Byte: 2}), p)
-	p, ok = normal.Schemas[0].Tables[0].Columns[0].Pos()
-	require.True(t, ok)
+	p = normal.Schemas[0].Tables[0].Columns[0].Pos()
 	require.Equal(t, schema.NewFilePos("schema.hcl").SetStart(hcl.Pos{Line: 3, Column: 3, Byte: 3}), p)
 }