Repository: ariga/atlas
Pull request #3197 (2024-10-20 11:45:21)

sql/schema: allow schema loaders to attach positions


Patch:
diff --git a/schemahcl/extension.go b/schemahcl/extension.go
index 0ca5ee0f..e43fc6dc 100644
--- a/schemahcl/extension.go
+++ b/schemahcl/extension.go
@@ -11,6 +11,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/hashicorp/hcl/v2"
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/gocty"
 )
@@ -96,10 +97,12 @@ func (r *Resource) As(target any) error {
 	return r.as(target)
 }
 
+var typeRange = reflect.TypeOf(&hcl.Range{})
+
 // As reads the attributes and children resources of the resource into the target struct.
 func (r *Resource) as(target any) error {
 	existingAttrs, existingChildren := existingElements(r)
-	var seenName, seenQualifier bool
+	var seenName, seenQualifier, setRange bool
 	v := reflect.ValueOf(target).Elem()
 	for _, ft := range specFields(target) {
 		field := v.FieldByName(ft.Name)
@@ -119,6 +122,14 @@ func (r *Resource) as(target any) error {
 			}
 			seenQualifier = true
 			field.SetString(r.Qualifier)
+		case ft.isRange() && !hasAttr(r, ft.tag):
+			if field.Type() != typeRange {
+				return fmt.Errorf("schemahcl: expected field %q to be of type *hcl.Range: %v", ft.Name, field.Type())
+			}
+			if r.rang != nil {
+				setRange = true
+				field.Set(reflect.ValueOf(r.rang))
+			}
 		case hasAttr(r, ft.tag):
 			attr, _ := r.Attr(ft.tag)
 			if err := setField(field, attr); err != nil {
@@ -221,9 +232,9 @@ func (r *Resource) as(target any) error {
 		children := childrenOfType(r, childType)
 		extras.Children = append(extras.Children, children...)
 	}
-	// In case the resource contains a remain (DefaultExtension)
-	// field, attach to it the position.
-	if r.rang != nil && rem.Remain() != nil {
+	// In case the resource contains a remain (DefaultExtension) and
+	// the range was not explicitly set, attach to it the position.
+	if r.rang != nil && rem.Remain() != nil && !setRange {
 		rem.Remain().SetRange(r.rang)
 	}
 	return nil
@@ -613,6 +624,8 @@ func (f fieldDesc) isName() bool { return f.is("name") }
 
 func (f fieldDesc) isQualifier() bool { return f.is("qualifier") }
 
+func (f fieldDesc) isRange() bool { return f.is("range") }
+
 func (f fieldDesc) omitempty() bool { return f.is("omitempty") }
 
 func (f fieldDesc) is(t string) bool {
diff --git a/schemahcl/spec.go b/schemahcl/spec.go
index 54fd5629..eac3a973 100644
--- a/schemahcl/spec.go
+++ b/schemahcl/spec.go
@@ -609,3 +609,12 @@ func StringEnumsAttr(k string, elems ...*EnumString) *Attr {
 		V: cty.ListVal(vv),
 	}
 }
+
+// RangeAsPos builds a schema position from the give HCL range.
+func RangeAsPos(r hcl.Range) *schema.Pos {
+	return &schema.Pos{
+		Filename: r.Filename,
+		Start:    r.Start,
+		End:      r.End,
+	}
+}
diff --git a/sql/schema/schema.go b/sql/schema/schema.go
index c0e96076..df2abda9 100644
--- a/sql/schema/schema.go
+++ b/sql/schema/schema.go
@@ -614,6 +614,17 @@ type (
 	Materialized struct {
 		Attr
 	}
+
+	// Pos is an attribute that holds the position of a schema element.
+	Pos struct {
+		// Filename is the name (or full path) of the file which loaded the schema element.
+		Filename string
+
+		// Start and End represent the bounds of this range.
+		Start, End struct {
+			Line, Column, Byte int // hcl.Pos fields.
+		}
+	}
 )
 
 // A list of known view check options.
@@ -650,6 +661,7 @@ func (*DecimalType) typ()     {}
 func (*UnsupportedType) typ() {}
 
 // attributes.
+func (*Pos) attr()             {}
 func (*Check) attr()           {}
 func (*Comment) attr()         {}
 func (*Charset) attr()         {}
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index 3b1b9893..faee6a77 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -10,6 +10,7 @@ import (
 
 	"ariga.io/atlas/schemahcl"
 
+	"github.com/hashicorp/hcl/v2"
 	"github.com/zclconf/go-cty/cty"
 )
 
@@ -31,6 +32,7 @@ type (
 		Indexes     []*Index       `spec:"index"`
 		Checks      []*Check       `spec:"check"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// View holds a specification for an SQL view.
@@ -45,6 +47,7 @@ type (
 		// The definition is appended as additional attribute
 		// by the spec creator to marshal it after the columns.
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// Column holds a specification for a column in an SQL table.
@@ -54,6 +57,7 @@ type (
 		Type    *schemahcl.Type `spec:"type"`
 		Default cty.Value       `spec:"default"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// PrimaryKey holds a specification for the primary key of a table.
@@ -61,6 +65,7 @@ type (
 		Parts   []*IndexPart     `spec:"on"`
 		Columns []*schemahcl.Ref `spec:"columns"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// Index holds a specification for the index key of a table.
@@ -70,6 +75,7 @@ type (
 		Parts   []*IndexPart     `spec:"on"`
 		Columns []*schemahcl.Ref `spec:"columns"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// IndexPart holds a specification for the index key part.
@@ -78,6 +84,7 @@ type (
 		Column *schemahcl.Ref `spec:"column"`
 		Expr   string         `spec:"expr,omitempty"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// Check holds a specification for a check constraint on a table.
@@ -85,6 +92,7 @@ type (
 		Name string `spec:",name"`
 		Expr string `spec:"expr"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// ForeignKey holds a specification for the Foreign key of a table.
@@ -95,6 +103,7 @@ type (
 		OnUpdate   *schemahcl.Ref   `spec:"on_update"`
 		OnDelete   *schemahcl.Ref   `spec:"on_delete"`
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// Func holds the specification for a function.
@@ -107,6 +116,7 @@ type (
 		// The definition and the return type are appended as additional
 		// attribute by the spec creator to marshal it after the arguments.
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// FuncArg holds the specification for a function argument.
@@ -117,6 +127,7 @@ type (
 		// Optional attributes such as mode are added by the driver,
 		// as their definition can be either a string or an enum (ref).
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// Trigger holds the specification for a trigger.
@@ -125,6 +136,7 @@ type (
 		On   *schemahcl.Ref `spec:"on"` // A table or a view.
 		// Attributes and blocks are different for each driver.
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 
 	// Sequence holds a specification for a Sequence.
@@ -135,6 +147,7 @@ type (
 		// Type, Start, Increment, Min, Max, Cache, Cycle
 		// are optionally added to the sequence definition.
 		schemahcl.DefaultExtension
+		Range *hcl.Range `spec:",range"`
 	}
 )
 


Test patch:
diff --git a/schemahcl/schemahcl_test.go b/schemahcl/schemahcl_test.go
index f2b58ae1..7f83873e 100644
--- a/schemahcl/schemahcl_test.go
+++ b/schemahcl/schemahcl_test.go
@@ -1138,6 +1138,11 @@ func Test_WithPos(t *testing.T) {
 				A int `spec:"a"`
 				DefaultExtension
 			} `spec:"bar"`
+			Qux struct {
+				Range *hcl.Range `spec:",range"`
+				A     int        `spec:"a"`
+				DefaultExtension
+			} `spec:"qux"`
 			DefaultExtension
 		}
 		b = []byte(`
@@ -1151,6 +1156,9 @@ baz = 1
 bar {
   a = 1
 }
+qux {
+  a = 1
+}
 `)
 	)
 	require.NoError(t, New(WithPos()).EvalBytes(b, &doc, nil))
@@ -1164,6 +1172,8 @@ bar {
 	require.Equal(t, 4, rs[1].Children[0].Range().Start.Line)
 	require.Equal(t, 5, rs[1].Children[0].Attrs[0].Range().Start.Line)
 	require.NotNil(t, doc.Bar.Extra.Range(), "position should be attached to the resource")
+	require.Equal(t, doc.Qux.Range.Start.Line, 12)
+	require.Nil(t, doc.Qux.Extra.Range(), "position should not be attached if it was explicitly set")
 }
 
 func TestExtendedBlockDef(t *testing.T) {
diff --git a/schemahcl/spec_test.go b/schemahcl/spec_test.go
index 79f30f99..97a67641 100644
--- a/schemahcl/spec_test.go
+++ b/schemahcl/spec_test.go
@@ -8,6 +8,9 @@ import (
 	"strconv"
 	"testing"
 
+	"ariga.io/atlas/sql/schema"
+
+	"github.com/hashicorp/hcl/v2"
 	"github.com/stretchr/testify/require"
 )
 
@@ -148,3 +151,17 @@ func TestBuildRef(t *testing.T) {
 		})
 	}
 }
+
+func TestRangeAsPos(t *testing.T) {
+	p := RangeAsPos(hcl.Range{
+		Filename: "schema.pg.hcl",
+	})
+	require.Equal(t, &schema.Pos{Filename: "schema.pg.hcl"}, p)
+
+	p = RangeAsPos(hcl.Range{
+		Filename: "schema.pg.hcl",
+		Start:    hcl.Pos{Byte: 10},
+		End:      hcl.Pos{Byte: 20},
+	})
+	require.Equal(t, &schema.Pos{Filename: "schema.pg.hcl", Start: hcl.Pos{Byte: 10}, End: hcl.Pos{Byte: 20}}, p)
+}