Repository: ariga/atlas
Pull request #1081 (2022-08-21 20:17:18)

sql/internal/specutil: patch foreign-key qualified references
The last part is to handle qualified enums as well.


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 46d19206..07f0f288 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -86,7 +86,7 @@ func findTableSpec(tableSpecs []*sqlspec.Table, schemaName, tableName string) (*
 // ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
 // at this point. Instead, the linking is done by the Schema function.
 func Table(spec *sqlspec.Table, parent *schema.Schema, convertColumn ConvertColumnFunc,
-	convertPk ConvertPrimaryKeyFunc, convertIndex ConvertIndexFunc, convertCheck ConvertCheckFunc) (*schema.Table, error) {
+	convertPK ConvertPrimaryKeyFunc, convertIndex ConvertIndexFunc, convertCheck ConvertCheckFunc) (*schema.Table, error) {
 	tbl := &schema.Table{
 		Name:   spec.Name,
 		Schema: parent,
@@ -99,7 +99,7 @@ func Table(spec *sqlspec.Table, parent *schema.Schema, convertColumn ConvertColu
 		tbl.Columns = append(tbl.Columns, col)
 	}
 	if spec.PrimaryKey != nil {
-		pk, err := convertPk(spec.PrimaryKey, tbl)
+		pk, err := convertPK(spec.PrimaryKey, tbl)
 		if err != nil {
 			return nil, err
 		}
@@ -506,11 +506,7 @@ func FromForeignKey(s *schema.ForeignKey) (*sqlspec.ForeignKey, error) {
 	for _, v := range s.RefColumns {
 		ref := ColumnRef(v.Name)
 		if s.Table != s.RefTable {
-			if s.Table.Schema == s.RefTable.Schema {
-				ref = externalColRef(v.Name, s.RefTable.Name)
-			} else {
-				ref = crossSchemaColRef(v.Name, s.RefTable.Name, s.RefTable.Schema.Name)
-			}
+			ref = externalColRef(v.Name, s.RefTable.Name)
 		}
 		r = append(r, ref)
 	}
@@ -635,7 +631,7 @@ func externalColRef(cName string, tName string) *schemahcl.Ref {
 	return &schemahcl.Ref{V: "$table." + tName + ".$column." + cName}
 }
 
-func crossSchemaColRef(cName, tName, sName string) *schemahcl.Ref {
+func qualifiedExternalColRef(cName, tName, sName string) *schemahcl.Ref {
 	return &schemahcl.Ref{V: "$table." + sName + "." + tName + ".$column." + cName}
 }
 
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index c059c254..1914d5d0 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -113,15 +113,15 @@ func Marshal(v any, marshaler schemahcl.Marshaler, schemaSpec func(schem *schema
 			d.Tables = append(d.Tables, tables...)
 			d.Schemas = append(d.Schemas, spec)
 		}
-		if err := QualifyReferencedTables(d.Tables, s); err != nil {
+		if err := QualifyDuplicates(d.Tables); err != nil {
+			return nil, err
+		}
+		if err := QualifyReferences(d.Tables, s); err != nil {
 			return nil, err
 		}
 	default:
 		return nil, fmt.Errorf("specutil: failed marshaling spec. %T is not supported", v)
 	}
-	if err := QualifyDuplicates(d.Tables); err != nil {
-		return nil, err
-	}
 	return marshaler.MarshalSpec(d)
 }
 
@@ -147,24 +147,42 @@ func QualifyDuplicates(tableSpecs []*sqlspec.Table) error {
 	return nil
 }
 
-// QualifyReferencedTables sets the Qualified field equal to the schema name in any tables
-// that are references across schema boundaries from any foreign key.
-func QualifyReferencedTables(tableSpecs []*sqlspec.Table, realm *schema.Realm) error {
-	// Iterate over all tables in the realm. If any of the tables has a foreign key referencing a table in another
-	// schema, the targeted table needs to have a qualifier set.
-	for _, s := range realm.Schemas {
-		for _, t := range s.Tables {
-			for _, f := range t.ForeignKeys {
-				if f.Table.Schema != f.RefTable.Schema {
-					for _, tbl := range tableSpecs {
-						n, err := SchemaName(tbl.Schema)
-						if err != nil {
-							return err
-						}
-						if f.RefTable.Name == tbl.Name && f.RefTable.Schema.Name == n && tbl.Qualifier == "" {
-							tbl.Qualifier = n
-						}
-					}
+// QualifyReferences qualifies any reference with qualifier.
+func QualifyReferences(tableSpecs []*sqlspec.Table, realm *schema.Realm) error {
+	type cref struct{ s, t string }
+	byRef := make(map[cref]*sqlspec.Table)
+	for _, t := range tableSpecs {
+		r := cref{s: t.Qualifier, t: t.Name}
+		if byRef[r] != nil {
+			return fmt.Errorf("duplicate references were found for: %v", r)
+		}
+		byRef[r] = t
+	}
+	for _, t := range tableSpecs {
+		sname, err := SchemaName(t.Schema)
+		if err != nil {
+			return err
+		}
+		s1, ok := realm.Schema(sname)
+		if !ok {
+			return fmt.Errorf("schema %q was not found in realm", sname)
+		}
+		t1, ok := s1.Table(t.Name)
+		if !ok {
+			return fmt.Errorf("table %q.%q was not found in realm", sname, t.Name)
+		}
+		for _, fk := range t.ForeignKeys {
+			fk1, ok := t1.ForeignKey(fk.Symbol)
+			if !ok {
+				return fmt.Errorf("table %q.%q.%q was not found in realm", sname, t.Name, fk.Symbol)
+			}
+			for i, c := range fk.RefColumns {
+				if r, ok := byRef[cref{s: fk1.RefTable.Schema.Name, t: fk1.RefTable.Name}]; ok && r.Qualifier != "" {
+					fk.RefColumns[i] = qualifiedExternalColRef(fk1.RefColumns[i].Name, r.Name, r.Qualifier)
+				} else if r, ok := byRef[cref{t: fk1.RefTable.Name}]; ok && r.Qualifier == "" {
+					fk.RefColumns[i] = externalColRef(fk1.RefColumns[i].Name, r.Name)
+				} else {
+					return fmt.Errorf("missing reference for column %q in %q.%q.%q", c.V, sname, t.Name, fk.Symbol)
 				}
 			}
 		}
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 1b030e41..37e04fec 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -94,15 +94,15 @@ func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
 			d.Schemas = append(d.Schemas, doc.Schemas...)
 			d.Enums = append(d.Enums, doc.Enums...)
 		}
-		if err := specutil.QualifyReferencedTables(d.Tables, s); err != nil {
+		if err := specutil.QualifyDuplicates(d.Tables); err != nil {
+			return nil, err
+		}
+		if err := specutil.QualifyReferences(d.Tables, s); err != nil {
 			return nil, err
 		}
 	default:
 		return nil, fmt.Errorf("specutil: failed marshaling spec. %T is not supported", v)
 	}
-	if err := specutil.QualifyDuplicates(d.Tables); err != nil {
-		return nil, err
-	}
 	return marshaler.MarshalSpec(&d)
 }
 


Test patch:
diff --git a/internal/integration/mysql_test.go b/internal/integration/mysql_test.go
index f73d8630..d2db1a48 100644
--- a/internal/integration/mysql_test.go
+++ b/internal/integration/mysql_test.go
@@ -742,6 +742,13 @@ func TestMySQL_HCL_ForeignKeyCrossSchema(t *testing.T) {
     columns = [column.user_id]
   }
 }
+table "financial" "users" {
+  schema = schema.financial
+  column "id" {
+    null = false
+    type = int
+  }
+}
 table "users" "users" {
   schema = schema.users
   column "id" {
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index 34c4a0ac..5488e6be 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -557,6 +557,13 @@ func TestPostgres_HCL_ForeignKeyCrossSchema(t *testing.T) {
     on_delete   = NO_ACTION
   }
 }
+table "financial" "users" {
+  schema = schema.financial
+  column "id" {
+    null = false
+    type = serial
+  }
+}
 table "users" "users" {
   schema = schema.users
   column "id" {
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index b24a30b4..19fd7c22 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -8,11 +8,8 @@ import (
 	"fmt"
 	"testing"
 
-	"ariga.io/atlas/schemahcl"
 	"ariga.io/atlas/sql/internal/spectest"
-	"ariga.io/atlas/sql/internal/specutil"
 	"ariga.io/atlas/sql/schema"
-	"ariga.io/atlas/sql/sqlspec"
 	"github.com/stretchr/testify/require"
 )
 
@@ -1189,41 +1186,6 @@ func TestInputVars(t *testing.T) {
 	spectest.TestInputVars(t, EvalHCL)
 }
 
-func TestQualifyReferencedTables(t *testing.T) {
-	var (
-		col       = schema.NewIntColumn("col", "integer")
-		targetCol = schema.NewIntColumn("target_column", "integer")
-		targetTbl = schema.NewTable("target_table").AddColumns(targetCol)
-		realm     = schema.NewRealm(
-			schema.New("target_schema").AddTables(targetTbl),
-			schema.New("sch").
-				AddTables(
-					schema.NewTable("tbl").
-						AddColumns(col).
-						AddForeignKeys(
-							schema.NewForeignKey("col_fk").
-								SetRefTable(targetTbl).
-								AddColumns(col).
-								AddRefColumns(targetCol),
-						),
-				),
-		)
-		tables = []*sqlspec.Table{
-			{
-				Name:   "tbl",
-				Schema: &schemahcl.Ref{V: "$schema.sch"},
-			},
-			{
-				Name:   "target_table",
-				Schema: &schemahcl.Ref{V: "$schema.target_schema"},
-			},
-		}
-	)
-	require.NoError(t, specutil.QualifyReferencedTables(tables, realm))
-	require.Zero(t, tables[0].Qualifier)
-	require.Equal(t, "target_schema", tables[1].Qualifier)
-}
-
 func TestParseType_Decimal(t *testing.T) {
 	for _, tt := range []struct {
 		input   string
@@ -1317,3 +1279,92 @@ table "s2" "t2" {
 		AddRefColumns(expected.Schemas[1].Tables[0].Columns[0]))
 	require.NoError(t, EvalHCLBytes(s, &r, nil))
 }
+
+func TestMarshalRealm(t *testing.T) {
+	t1 := schema.NewTable("t1").
+		AddColumns(schema.NewIntColumn("id", "int"))
+	t2 := schema.NewTable("t2").
+		SetComment("Qualified with s1").
+		AddColumns(schema.NewIntColumn("oid", "int"))
+	t2.AddForeignKeys(schema.NewForeignKey("oid2id").AddColumns(t2.Columns[0]).SetRefTable(t1).AddRefColumns(t1.Columns[0]))
+
+	t3 := schema.NewTable("t3").
+		AddColumns(schema.NewIntColumn("id", "int"))
+	t4 := schema.NewTable("t2").
+		SetComment("Qualified with s2").
+		AddColumns(schema.NewIntColumn("oid", "int"))
+	t4.AddForeignKeys(schema.NewForeignKey("oid2id").AddColumns(t4.Columns[0]).SetRefTable(t3).AddRefColumns(t3.Columns[0]))
+	t5 := schema.NewTable("t5").
+		AddColumns(schema.NewIntColumn("oid", "int"))
+	t5.AddForeignKeys(schema.NewForeignKey("oid2id1").AddColumns(t5.Columns[0]).SetRefTable(t1).AddRefColumns(t1.Columns[0]))
+	// Reference is qualified with s1.
+	t5.AddForeignKeys(schema.NewForeignKey("oid2id2").AddColumns(t5.Columns[0]).SetRefTable(t2).AddRefColumns(t2.Columns[0]))
+
+	r := schema.NewRealm(
+		schema.New("s1").AddTables(t1, t2),
+		schema.New("s2").AddTables(t3, t4, t5),
+	)
+	got, err := MarshalHCL.MarshalSpec(r)
+	require.NoError(t, err)
+	require.Equal(
+		t,
+		`table "t1" {
+  schema = schema.s1
+  column "id" {
+    null = false
+    type = int
+  }
+}
+table "s1" "t2" {
+  schema  = schema.s1
+  comment = "Qualified with s1"
+  column "oid" {
+    null = false
+    type = int
+  }
+  foreign_key "oid2id" {
+    columns     = [column.oid]
+    ref_columns = [table.t1.column.id]
+  }
+}
+table "t3" {
+  schema = schema.s2
+  column "id" {
+    null = false
+    type = int
+  }
+}
+table "s2" "t2" {
+  schema  = schema.s2
+  comment = "Qualified with s2"
+  column "oid" {
+    null = false
+    type = int
+  }
+  foreign_key "oid2id" {
+    columns     = [column.oid]
+    ref_columns = [table.t3.column.id]
+  }
+}
+table "t5" {
+  schema = schema.s2
+  column "oid" {
+    null = false
+    type = int
+  }
+  foreign_key "oid2id1" {
+    columns     = [column.oid]
+    ref_columns = [table.t1.column.id]
+  }
+  foreign_key "oid2id2" {
+    columns     = [column.oid]
+    ref_columns = [table.s1.t2.column.oid]
+  }
+}
+schema "s1" {
+}
+schema "s2" {
+}
+`,
+		string(got))
+}
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index bb72cdbb..8c8a8072 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -9,11 +9,8 @@ import (
 	"strconv"
 	"testing"
 
-	"ariga.io/atlas/schemahcl"
 	"ariga.io/atlas/sql/internal/spectest"
-	"ariga.io/atlas/sql/internal/specutil"
 	"ariga.io/atlas/sql/schema"
-	"ariga.io/atlas/sql/sqlspec"
 	"github.com/stretchr/testify/require"
 )
 
@@ -1284,37 +1281,91 @@ func TestInputVars(t *testing.T) {
 	spectest.TestInputVars(t, EvalHCL)
 }
 
-func TestQualifyReferencedTables(t *testing.T) {
-	var (
-		col       = schema.NewIntColumn("col", "integer")
-		targetCol = schema.NewIntColumn("target_column", "integer")
-		targetTbl = schema.NewTable("target_table").AddColumns(targetCol)
-		realm     = schema.NewRealm(
-			schema.New("target_schema").AddTables(targetTbl),
-			schema.New("sch").
-				AddTables(
-					schema.NewTable("tbl").
-						AddColumns(col).
-						AddForeignKeys(
-							schema.NewForeignKey("col_fk").
-								SetRefTable(targetTbl).
-								AddColumns(col).
-								AddRefColumns(targetCol),
-						),
-				),
-		)
-		tables = []*sqlspec.Table{
-			{
-				Name:   "tbl",
-				Schema: &schemahcl.Ref{V: "$schema.sch"},
-			},
-			{
-				Name:   "target_table",
-				Schema: &schemahcl.Ref{V: "$schema.target_schema"},
-			},
-		}
+func TestMarshalRealm(t *testing.T) {
+	t1 := schema.NewTable("t1").
+		AddColumns(schema.NewIntColumn("id", "int"))
+	t2 := schema.NewTable("t2").
+		SetComment("Qualified with s1").
+		AddColumns(schema.NewIntColumn("oid", "int"))
+	t2.AddForeignKeys(schema.NewForeignKey("oid2id").AddColumns(t2.Columns[0]).SetRefTable(t1).AddRefColumns(t1.Columns[0]))
+
+	t3 := schema.NewTable("t3").
+		AddColumns(schema.NewIntColumn("id", "int"))
+	t4 := schema.NewTable("t2").
+		SetComment("Qualified with s2").
+		AddColumns(schema.NewIntColumn("oid", "int"))
+	t4.AddForeignKeys(schema.NewForeignKey("oid2id").AddColumns(t4.Columns[0]).SetRefTable(t3).AddRefColumns(t3.Columns[0]))
+	t5 := schema.NewTable("t5").
+		AddColumns(schema.NewIntColumn("oid", "int"))
+	t5.AddForeignKeys(schema.NewForeignKey("oid2id1").AddColumns(t5.Columns[0]).SetRefTable(t1).AddRefColumns(t1.Columns[0]))
+	// Reference is qualified with s1.
+	t5.AddForeignKeys(schema.NewForeignKey("oid2id2").AddColumns(t5.Columns[0]).SetRefTable(t2).AddRefColumns(t2.Columns[0]))
+
+	r := schema.NewRealm(
+		schema.New("s1").AddTables(t1, t2),
+		schema.New("s2").AddTables(t3, t4, t5),
 	)
-	require.NoError(t, specutil.QualifyReferencedTables(tables, realm))
-	require.Zero(t, tables[0].Qualifier)
-	require.Equal(t, "target_schema", tables[1].Qualifier)
+	got, err := MarshalHCL.MarshalSpec(r)
+	require.NoError(t, err)
+	require.Equal(
+		t,
+		`table "t1" {
+  schema = schema.s1
+  column "id" {
+    null = false
+    type = int
+  }
+}
+table "s1" "t2" {
+  schema  = schema.s1
+  comment = "Qualified with s1"
+  column "oid" {
+    null = false
+    type = int
+  }
+  foreign_key "oid2id" {
+    columns     = [column.oid]
+    ref_columns = [table.t1.column.id]
+  }
+}
+table "t3" {
+  schema = schema.s2
+  column "id" {
+    null = false
+    type = int
+  }
+}
+table "s2" "t2" {
+  schema  = schema.s2
+  comment = "Qualified with s2"
+  column "oid" {
+    null = false
+    type = int
+  }
+  foreign_key "oid2id" {
+    columns     = [column.oid]
+    ref_columns = [table.t3.column.id]
+  }
+}
+table "t5" {
+  schema = schema.s2
+  column "oid" {
+    null = false
+    type = int
+  }
+  foreign_key "oid2id1" {
+    columns     = [column.oid]
+    ref_columns = [table.t1.column.id]
+  }
+  foreign_key "oid2id2" {
+    columns     = [column.oid]
+    ref_columns = [table.s1.t2.column.oid]
+  }
+}
+schema "s1" {
+}
+schema "s2" {
+}
+`,
+		string(got))
 }