Repository: ariga/atlas
Pull request #2005 (2023-08-18 13:38:59)

sql: avoid ambiguity between resource label and qualifier
Resolve: https://github.com/ariga/atlas/issues/1984


Patch:
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index 203df7e3..fdca87b7 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -58,10 +58,10 @@ func Marshal(v any, marshaler schemahcl.Marshaler, convertFunc func(*schema.Sche
 			d.Views = append(d.Views, spec.Views...)
 			d.Schemas = append(d.Schemas, spec.Schema)
 		}
-		if err := QualifyTables(d.Tables); err != nil {
+		if err := QualifyObjects(d.Tables); err != nil {
 			return nil, err
 		}
-		if err := QualifyViews(d.Views); err != nil {
+		if err := QualifyObjects(d.Views); err != nil {
 			return nil, err
 		}
 		if err := QualifyReferences(d.Tables, s); err != nil {
@@ -73,46 +73,53 @@ func Marshal(v any, marshaler schemahcl.Marshaler, convertFunc func(*schema.Sche
 	return marshaler.MarshalSpec(d)
 }
 
-// QualifyTables sets the Qualifier field equal to the schema
-// name in any tables with duplicate names in the provided specs.
-func QualifyTables(specs []*sqlspec.Table) error {
-	seen := make(map[string]*sqlspec.Table, len(specs))
-	for _, t := range specs {
-		if s, ok := seen[t.Name]; ok {
-			schemaName, err := SchemaName(s.Schema)
+// SchemaObject describes a top-level schema object
+// that might be qualified, e.g. a table or a view.
+type SchemaObject interface {
+	Label() string
+	QualifierLabel() string
+	SetQualifier(string)
+	SchemaRef() *schemahcl.Ref
+}
+
+// QualifyObjects sets the Qualifier field equal to the schema
+// name in any objects with duplicate names in the provided specs.
+func QualifyObjects[T SchemaObject](specs []T) error {
+	var (
+		seen    = make(map[string]T, len(specs))
+		schemas = make(map[string]bool, len(specs))
+	)
+	// Loop first and qualify schema objects with the same label.
+	// For example, two tables named "users" reside in different
+	// schemas are converted to: ("s1", "users") and ("s2", "users").
+	for _, v := range specs {
+		if s, ok := seen[v.Label()]; ok {
+			schemaName, err := SchemaName(s.SchemaRef())
 			if err != nil {
 				return err
 			}
-			s.Qualifier = schemaName
-			schemaName, err = SchemaName(t.Schema)
+			s.SetQualifier(schemaName)
+			schemas[schemaName] = true
+			schemaName, err = SchemaName(v.SchemaRef())
 			if err != nil {
 				return err
 			}
-			t.Qualifier = schemaName
+			v.SetQualifier(schemaName)
+			schemas[schemaName] = true
 		}
-		seen[t.Name] = t
+		seen[v.Label()] = v
 	}
-	return nil
-}
-
-// QualifyViews sets the Qualifier field equal to the schema
-// name in any tables with duplicate names in the provided specs.
-func QualifyViews(specs []*sqlspec.View) error {
-	seen := make(map[string]*sqlspec.View, len(specs))
+	// After objects were qualified, they might be conflicted with different
+	// resources that labeled with the schema name. e.g., ("s1", "users") and
+	// ("s1"). To resolve this conflict, we qualify these objects as well.
 	for _, v := range specs {
-		if s, ok := seen[v.Name]; ok {
-			schemaName, err := SchemaName(s.Schema)
-			if err != nil {
-				return err
-			}
-			s.Qualifier = schemaName
-			schemaName, err = SchemaName(v.Schema)
+		if v.QualifierLabel() == "" && schemas[v.Label()] {
+			schemaName, err := SchemaName(v.SchemaRef())
 			if err != nil {
 				return err
 			}
-			v.Qualifier = schemaName
+			v.SetQualifier(schemaName)
 		}
-		seen[v.Name] = v
 	}
 	return nil
 }
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 938362ab..a57590f1 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -30,13 +30,26 @@ type (
 	}
 	// Enum holds a specification for an enum, that can be referenced as a column type.
 	Enum struct {
-		Name   string         `spec:",name"`
-		Schema *schemahcl.Ref `spec:"schema"`
-		Values []string       `spec:"values"`
+		Name      string         `spec:",name"`
+		Qualifier string         `spec:",qualifier"`
+		Schema    *schemahcl.Ref `spec:"schema"`
+		Values    []string       `spec:"values"`
 		schemahcl.DefaultExtension
 	}
 )
 
+// Label returns the defaults label used for the enum resource.
+func (e *Enum) Label() string { return e.Name }
+
+// QualifierLabel returns the qualifier label used for the enum resource, if any.
+func (e *Enum) QualifierLabel() string { return e.Qualifier }
+
+// SetQualifier sets the qualifier label used for the enum resource.
+func (e *Enum) SetQualifier(q string) { e.Qualifier = q }
+
+// SchemaRef returns the schema reference for the enum.
+func (e *Enum) SchemaRef() *schemahcl.Ref { return e.Schema }
+
 func init() {
 	schemahcl.Register("enum", &Enum{})
 }
@@ -112,10 +125,13 @@ 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.QualifyTables(d.Tables); err != nil {
+		if err := specutil.QualifyObjects(d.Tables); err != nil {
+			return nil, err
+		}
+		if err := specutil.QualifyObjects(d.Views); err != nil {
 			return nil, err
 		}
-		if err := specutil.QualifyViews(d.Views); err != nil {
+		if err := specutil.QualifyObjects(d.Enums); err != nil {
 			return nil, err
 		}
 		if err := specutil.QualifyReferences(d.Tables, s); err != nil {
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index d797fcaf..16ecc4a9 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -94,6 +94,30 @@ type (
 	Type string
 )
 
+// Label returns the defaults label used for the table resource.
+func (t *Table) Label() string { return t.Name }
+
+// QualifierLabel returns the qualifier label used for the table resource, if any.
+func (t *Table) QualifierLabel() string { return t.Qualifier }
+
+// SetQualifier sets the qualifier label used for the table resource.
+func (t *Table) SetQualifier(q string) { t.Qualifier = q }
+
+// SchemaRef returns the schema reference for the table.
+func (t *Table) SchemaRef() *schemahcl.Ref { return t.Schema }
+
+// Label returns the defaults label used for the view resource.
+func (v *View) Label() string { return v.Name }
+
+// QualifierLabel returns the qualifier label used for the table resource, if any.
+func (v *View) QualifierLabel() string { return v.Qualifier }
+
+// SetQualifier sets the qualifier label used for the view resource.
+func (v *View) SetQualifier(q string) { v.Qualifier = q }
+
+// SchemaRef returns the schema reference for the view.
+func (v *View) SchemaRef() *schemahcl.Ref { return v.Schema }
+
 func init() {
 	schemahcl.Register("view", &View{})
 	schemahcl.Register("table", &Table{})


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 56cb2408..d876eb5c 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -2116,6 +2116,53 @@ schema "s2" {
 		string(got))
 }
 
+func TestMarshalQualifiers(t *testing.T) {
+	var (
+		s1 = schema.New("s1").
+			AddTables(
+				schema.NewTable("t1").AddColumns(schema.NewIntColumn("id", "int")),
+			)
+		s2 = schema.New("s2").
+			AddTables(
+				schema.NewTable("t1").AddColumns(schema.NewIntColumn("id", "int")),
+			)
+		s3 = schema.New("s3").
+			AddTables(
+				schema.NewTable("s1").AddColumns(schema.NewIntColumn("id", "int")),
+			)
+	)
+	buf, err := MarshalHCL.MarshalSpec(schema.NewRealm(s1, s2, s3))
+	require.NoError(t, err)
+	require.Equal(t, `table "s1" "t1" {
+  schema = schema.s1
+  column "id" {
+    null = false
+    type = int
+  }
+}
+table "s2" "t1" {
+  schema = schema.s2
+  column "id" {
+    null = false
+    type = int
+  }
+}
+table "s3" "s1" {
+  schema = schema.s3
+  column "id" {
+    null = false
+    type = int
+  }
+}
+schema "s1" {
+}
+schema "s2" {
+}
+schema "s3" {
+}
+`, string(buf))
+}
+
 func TestEvalHCLError(t *testing.T) {
 	for _, tt := range []struct {
 		text, err string