Repository: ariga/atlas
Pull request #2972 (2024-07-16 18:07:17)

sql/schema: support exclude any spec type


Patch:
diff --git a/sql/schema/exclude.go b/sql/schema/exclude.go
index 983de581..2d0604bb 100644
--- a/sql/schema/exclude.go
+++ b/sql/schema/exclude.go
@@ -88,25 +88,41 @@ func split(patterns []string) ([][]string, error) {
 }
 
 func excludeS(s *Schema, glob []string) error {
-	if globE, exclude := excludeType(typeE, glob[0]); exclude {
-		var objects []Object
-		for _, t := range s.Objects {
-			e, ok := t.(*EnumType)
-			if !ok {
-				objects = append(objects, t)
-				continue
-			}
-			match, err := filepath.Match(globE, e.T)
+	var (
+		objects = make([]Object, 0, len(s.Objects))
+		t2glob  = make(map[string]struct {
+			glob    string
+			exclude bool
+		})
+	)
+	for _, o := range s.Objects {
+		nt, ok := o.(interface {
+			SpecType() string
+			SpecName() string
+		})
+		if !ok {
+			objects = append(objects, o)
+			continue
+		}
+		cache, ok := t2glob[nt.SpecType()]
+		if !ok {
+			cache.glob, cache.exclude = excludeType(nt.SpecType(), glob[0])
+			t2glob[nt.SpecType()] = cache
+		}
+		if cache.exclude {
+			match, err := filepath.Match(cache.glob, nt.SpecName())
 			if err != nil {
 				return err
 			}
 			// No match or glob has more than one pattern.
 			if !match || len(glob) != 1 {
-				objects = append(objects, t)
+				objects = append(objects, o)
 			}
+		} else {
+			objects = append(objects, o)
 		}
-		s.Objects = objects
 	}
+	s.Objects = objects
 	if globT, exclude := excludeType(typeT, glob[0]); exclude {
 		var tables []*Table
 		for _, t := range s.Tables {
@@ -255,10 +271,10 @@ const (
 	typePr = "procedure"
 )
 
-var reTypeVT = regexp.MustCompile(`\[type=([a-z|]+)+\]$`)
+var reType = regexp.MustCompile(`\[type=([a-z|]+)+\]$`)
 
 func excludeType(t, v string) (string, bool) {
-	matches := reTypeVT.FindStringSubmatch(v)
+	matches := reType.FindStringSubmatch(v)
 	if len(matches) != 2 {
 		return v, true
 	}


Test patch:
diff --git a/sql/schema/exclude_test.go b/sql/schema/exclude_test.go
index cb5006dc..7f516310 100644
--- a/sql/schema/exclude_test.go
+++ b/sql/schema/exclude_test.go
@@ -130,6 +130,83 @@ func TestExcludeRealm_Enums(t *testing.T) {
 	require.Empty(t, r.Schemas[2].Objects)
 }
 
+type (
+	domain struct {
+		Object
+		name string
+	}
+	composite struct {
+		Object
+		name string
+	}
+)
+
+func (*domain) SpecType() string   { return "domain" }
+func (d *domain) SpecName() string { return d.name }
+
+func (*composite) SpecType() string   { return "composite" }
+func (c *composite) SpecName() string { return c.name }
+
+func TestExcludeRealm_ExternalType(t *testing.T) {
+	e1, e2, e3 := &EnumType{T: "e1"}, &EnumType{T: "e2"}, &EnumType{T: "e3"}
+	d1, d2, d3 := &domain{name: "e1"}, &domain{name: "e2"}, &domain{name: "e3"}
+	c1, c2, c3 := &composite{name: "e1"}, &composite{name: "e2"}, &composite{name: "e3"}
+	s1, s2, s3 := New("s1").AddObjects(e1, d1, c1), New("s2").AddObjects(e1, e2, d1, d2, c1, c2), New("s3").AddObjects(e1, e2, e3, d1, d2, d3, c1, c2, c3)
+
+	r, err := ExcludeRealm(NewRealm(s1, s2, s3), []string{"s1.e1"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 3)
+	require.Empty(t, r.Schemas[0].Objects)
+	require.Len(t, r.Schemas[1].Objects, 6)
+	require.Len(t, r.Schemas[2].Objects, 9)
+
+	// Wrong selector.
+	r, err = ExcludeRealm(NewRealm(s1, s2, s3), []string{"*.e1[type=table]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 3)
+	require.Empty(t, r.Schemas[0].Objects)
+	require.Len(t, r.Schemas[1].Objects, 6)
+	require.Len(t, r.Schemas[2].Objects, 9)
+
+	// Enum selector.
+	r, err = ExcludeRealm(NewRealm(s1, s2, s3), []string{"*.e1[type=enum]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 3)
+	require.Empty(t, r.Schemas[0].Objects)
+	require.Len(t, r.Schemas[1].Objects, 5)
+	require.Equal(t, []Object{e2, d1, d2, c1, c2}, r.Schemas[1].Objects)
+	require.Len(t, r.Schemas[2].Objects, 8)
+	require.Equal(t, []Object{e2, e3, d1, d2, d3, c1, c2, c3}, r.Schemas[2].Objects)
+
+	// Domain and Composite selectors with name.
+	r, err = ExcludeRealm(NewRealm(s1, s2, s3), []string{"*.e2[type=domain|composite]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 3)
+	require.Empty(t, r.Schemas[0].Objects)
+	require.Len(t, r.Schemas[1].Objects, 3)
+	require.Equal(t, []Object{e2, d1, c1}, r.Schemas[1].Objects)
+	require.Len(t, r.Schemas[2].Objects, 6)
+	require.Equal(t, []Object{e2, e3, d1, d3, c1, c3}, r.Schemas[2].Objects)
+
+	// Domain and Composite selectors without name.
+	r, err = ExcludeRealm(NewRealm(s1, s2, s3), []string{"*.*[type=domain|composite]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 3)
+	require.Empty(t, r.Schemas[0].Objects)
+	require.Len(t, r.Schemas[1].Objects, 1)
+	require.Equal(t, []Object{e2}, r.Schemas[1].Objects)
+	require.Len(t, r.Schemas[2].Objects, 2)
+	require.Equal(t, []Object{e2, e3}, r.Schemas[2].Objects)
+
+	// Exclude all.
+	r, err = ExcludeRealm(NewRealm(s1, s2, s3), []string{"*.*"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 3)
+	require.Empty(t, r.Schemas[0].Objects)
+	require.Empty(t, r.Schemas[1].Objects)
+	require.Empty(t, r.Schemas[2].Objects)
+}
+
 func TestExcludeSchema_Enums(t *testing.T) {
 	e1, e2, e3 := &EnumType{T: "e1"}, &EnumType{T: "e2"}, &EnumType{T: "e3"}
 	r := NewRealm(New("s1").AddObjects(e1, e2, e3))