Repository: ariga/atlas
Pull request #2906 (2024-06-30 08:24:46)

sql/schema: add back-ref to schema dependencies


Patch:
diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
index 9f4a3b76..b6076cc0 100644
--- a/sql/schema/dsl.go
+++ b/sql/schema/dsl.go
@@ -253,9 +253,30 @@ func (t *Table) AddAttrs(attrs ...Attr) *Table {
 // AddDeps adds the given objects as dependencies to the view.
 func (t *Table) AddDeps(objs ...Object) *Table {
 	t.Deps = append(t.Deps, objs...)
+	addRefs(t, objs)
 	return t
 }
 
+// RefsAdder wraps the AddRefs method. Objects that implemented this method
+// will get their dependent objects automatically set by their AddDeps calls.
+type RefsAdder interface {
+	AddRefs(...Object)
+}
+
+// addRefs adds the dependent objects to all objects it references.
+func addRefs(dependent Object, refs []Object) {
+	for _, o := range refs {
+		if r, ok := o.(RefsAdder); ok {
+			r.AddRefs(dependent)
+		}
+	}
+}
+
+// AddRefs adds references to the table.
+func (t *Table) AddRefs(refs ...Object) {
+	t.Refs = append(t.Refs, refs...)
+}
+
 // NewView creates a new View.
 func NewView(name, def string) *View {
 	return &View{Name: name, Def: def}
@@ -295,9 +316,15 @@ func (v *View) AddAttrs(attrs ...Attr) *View {
 // AddDeps adds the given objects as dependencies to the view.
 func (v *View) AddDeps(objs ...Object) *View {
 	v.Deps = append(v.Deps, objs...)
+	addRefs(v, objs)
 	return v
 }
 
+// AddRefs adds references to the view.
+func (v *View) AddRefs(refs ...Object) {
+	v.Refs = append(v.Refs, refs...)
+}
+
 // AddIndexes appends the given indexes to the table index list.
 func (v *View) AddIndexes(indexes ...*Index) *View {
 	for _, idx := range indexes {
@@ -849,15 +876,27 @@ func (f *ForeignKey) SetOnDelete(o ReferenceOption) *ForeignKey {
 // AddDeps adds the given objects as dependencies to the function.
 func (f *Func) AddDeps(objs ...Object) *Func {
 	f.Deps = append(f.Deps, objs...)
+	addRefs(f, objs)
 	return f
 }
 
+// AddRefs adds references to the function.
+func (f *Func) AddRefs(refs ...Object) {
+	f.Refs = append(f.Refs, refs...)
+}
+
 // AddDeps adds the given objects as dependencies to the procedure.
 func (p *Proc) AddDeps(objs ...Object) *Proc {
 	p.Deps = append(p.Deps, objs...)
+	addRefs(p, objs)
 	return p
 }
 
+// AddRefs adds references to the procedure.
+func (p *Proc) AddRefs(refs ...Object) {
+	p.Refs = append(p.Refs, refs...)
+}
+
 // ReplaceOrAppend searches an attribute of the same type as v in
 // the list and replaces it. Otherwise, v is appended to the list.
 func ReplaceOrAppend(attrs *[]Attr, v Attr) {
diff --git a/sql/schema/schema.go b/sql/schema/schema.go
index ec49bc43..39f6aae9 100644
--- a/sql/schema/schema.go
+++ b/sql/schema/schema.go
@@ -43,6 +43,7 @@ type (
 		Attrs       []Attr     // Attrs, constraints and options.
 		Triggers    []*Trigger // Triggers on the table.
 		Deps        []Object   // Objects this table depends on.
+		Refs        []Object   // Objects that depends on this table.
 	}
 
 	// A View represents a view definition.
@@ -55,6 +56,7 @@ type (
 		Indexes  []*Index   // Indexes on materialized view.
 		Triggers []*Trigger // Triggers on the view.
 		Deps     []Object   // Objects this view depends on.
+		Refs     []Object   // Objects that depends on this view.
 	}
 
 	// A Column represents a column definition.
@@ -124,6 +126,7 @@ type (
 		Body       string         // Trigger body only.
 		Attrs      []Attr         // WHEN, REFERENCING, etc.
 		Deps       []Object       // Objects this trigger depends on.
+		Refs       []Object       // Objects that depend on this trigger.
 	}
 
 	// TriggerTime represents the trigger action time.
@@ -148,6 +151,7 @@ type (
 		Lang   string   // Language (e.g. SQL, PL/pgSQL, etc.).
 		Attrs  []Attr   // Extra driver specific attributes.
 		Deps   []Object // Objects this function depends on.
+		Refs   []Object // Objects that depend on this function.
 	}
 
 	// Proc represents a procedure definition.
@@ -159,6 +163,7 @@ type (
 		Lang   string   // Language (e.g. SQL, PL/pgSQL, etc.).
 		Attrs  []Attr   // Extra driver specific attributes.
 		Deps   []Object // Objects this function depends on.
+		Refs   []Object // Objects that depend on this Proc.
 	}
 
 	// A FuncArg represents a single function argument.
diff --git a/sql/sqlite/convert.go b/sql/sqlite/convert.go
index f5c18604..55ef8cb0 100644
--- a/sql/sqlite/convert.go
+++ b/sql/sqlite/convert.go
@@ -95,7 +95,7 @@ func ParseType(c string) (schema.Type, error) {
 			ct.Size = int(p)
 		}
 		return ct, nil
-	case "json":
+	case "json", "jsonb":
 		return &schema.JSONType{T: t}, nil
 	case "date", "datetime", "time", "timestamp":
 		return &schema.TimeType{T: t}, nil
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index 7581128e..5c469f58 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -280,6 +280,7 @@ var TypeRegistry = schemahcl.NewRegistry(
 		schemahcl.NewTypeSpec("datetime"),
 		schemahcl.NewTypeSpec("json"),
 		schemahcl.NewTypeSpec("uuid"),
+		schemahcl.NewTypeSpec("jsonb"),
 	),
 )
 


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index f4df7db6..52e9ea5f 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -522,6 +522,18 @@ schema "other" {}
 	)
 	r.AddSchemas(public, other)
 	require.NoError(t, EvalHCLBytes([]byte(f), &got, nil))
+	// References might be added in different order. However, since
+	// they are not marshaled, we just check that they were added.
+	for i, s := range got.Schemas {
+		for j, tr := range s.Tables {
+			require.Equal(t, len(tr.Refs), len(r.Schemas[i].Tables[j].Refs))
+			tr.Refs, r.Schemas[i].Tables[j].Refs = nil, nil
+		}
+		for j, v := range s.Views {
+			require.Equal(t, len(v.Refs), len(r.Schemas[i].Views[j].Refs))
+			v.Refs, r.Schemas[i].Views[j].Refs = nil, nil
+		}
+	}
 	require.EqualValues(t, r, got)
 }
 
diff --git a/sql/schema/dsl_test.go b/sql/schema/dsl_test.go
index e4b52ff9..ee039ad4 100644
--- a/sql/schema/dsl_test.go
+++ b/sql/schema/dsl_test.go
@@ -141,6 +141,9 @@ func TestSchema_Views(t *testing.T) {
 	v1, v2 := schema.NewView("v1", "SELECT 1"), schema.NewView("v2", "SELECT 2")
 	s.AddViews(v1, v2)
 	require.Equal(t, []*schema.View{v1, v2}, s.Views)
+	v1.AddDeps(v2)
+	require.Equal(t, []schema.Object{v2}, v1.Deps)
+	require.Equal(t, []schema.Object{v1}, v2.Refs)
 }
 
 func TestSchema_SetCharset(t *testing.T) {
diff --git a/sql/sqlite/inspect_test.go b/sql/sqlite/inspect_test.go
index f9f021eb..7d1db781 100644
--- a/sql/sqlite/inspect_test.go
+++ b/sql/sqlite/inspect_test.go
@@ -30,6 +30,7 @@ func TestDriver_InspectTable(t *testing.T) {
 					WillReturnRows(sqltest.Rows(`
  name |   type       | nullable | dflt_value  | primary  | hidden
 ------+--------------+----------+ ------------+----------+----------
+ id   | integer       |  0      |     0x1     |  1       |  0
  c1   | int           |  1      |     a       |  0       |  0
  c2   | integer       |  0      |     97      |  0       |  0
  c3   | varchar(100)  |  1      |    'A'      |  0       |  0
@@ -44,7 +45,7 @@ func TestDriver_InspectTable(t *testing.T) {
  w    | int           |  0      |             |  0       |  2
  x    | text          |  0      |             |  0       |  3
  y    | text          |  0      |             |  0       |  2
- id   | integer       |  0      |     0x1     |  1       |  0
+ j    | jsonb         |  0      |             |  0       |  0
 `))
 				m.noIndexes("users")
 				m.noFKs("users")
@@ -52,6 +53,7 @@ func TestDriver_InspectTable(t *testing.T) {
 			expect: func(require *require.Assertions, t *schema.Table, err error) {
 				require.NoError(err)
 				columns := []*schema.Column{
+					{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, Attrs: []schema.Attr{&AutoIncrement{}}, Default: &schema.Literal{V: "0x1"}},
 					{Name: "c1", Type: &schema.ColumnType{Null: true, Type: &schema.IntegerType{T: "int"}, Raw: "int"}, Default: &schema.RawExpr{X: "a"}},
 					{Name: "c2", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, Default: &schema.Literal{V: "97"}},
 					{Name: "c3", Type: &schema.ColumnType{Null: true, Type: &schema.StringType{T: "varchar", Size: 100}, Raw: "varchar(100)"}, Default: &schema.Literal{V: "'A'"}},
@@ -66,14 +68,14 @@ func TestDriver_InspectTable(t *testing.T) {
 					{Name: "w", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}, Raw: "int"}, Attrs: []schema.Attr{&schema.GeneratedExpr{Type: "VIRTUAL", Expr: "(a*10)"}}},
 					{Name: "x", Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}, Raw: "text"}, Attrs: []schema.Attr{&schema.GeneratedExpr{Type: "STORED", Expr: "(typeof(c))"}}},
 					{Name: "y", Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}, Raw: "text"}, Attrs: []schema.Attr{&schema.GeneratedExpr{Type: "VIRTUAL", Expr: "(substr(b,a,a+2))"}}},
-					{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, Attrs: []schema.Attr{&AutoIncrement{}}, Default: &schema.Literal{V: "0x1"}},
+					{Name: "j", Type: &schema.ColumnType{Type: &schema.JSONType{T: "jsonb"}, Raw: "jsonb"}},
 				}
 				require.Equal(t.Columns, columns)
 				require.EqualValues(&schema.Index{
 					Name:   "PRIMARY",
 					Unique: true,
 					Table:  t,
-					Parts:  []*schema.IndexPart{{SeqNo: 1, C: columns[len(columns)-1]}},
+					Parts:  []*schema.IndexPart{{SeqNo: 1, C: columns[0]}},
 					Attrs:  []schema.Attr{&AutoIncrement{}},
 				}, t.PrimaryKey)
 			},
diff --git a/sql/sqlite/sqlspec_test.go b/sql/sqlite/sqlspec_test.go
index 4b214552..8f83dfb5 100644
--- a/sql/sqlite/sqlspec_test.go
+++ b/sql/sqlite/sqlspec_test.go
@@ -458,6 +458,10 @@ func TestTypes(t *testing.T) {
 			typeExpr: "uuid",
 			expected: &schema.UUIDType{T: "uuid"},
 		},
+		{
+			typeExpr: "jsonb",
+			expected: &schema.JSONType{T: "jsonb"},
+		},
 	} {
 		t.Run(tt.typeExpr, func(t *testing.T) {
 			var test schema.Schema