Repository: ariga/atlas
Pull request #3230 (2024-11-24 19:28:30)

sql/schema: support excluding realm objects


Patch:
diff --git a/sql/schema/exclude.go b/sql/schema/exclude.go
index 6c9d4965..426cc416 100644
--- a/sql/schema/exclude.go
+++ b/sql/schema/exclude.go
@@ -17,30 +17,41 @@ func ExcludeRealm(r *Realm, patterns []string) (*Realm, error) {
 	if len(patterns) == 0 {
 		return r, nil
 	}
-	var schemas []*Schema
 	globs, err := split(patterns)
 	if err != nil {
 		return nil, err
 	}
+	for _, g := range globs {
+		// Realm objects are top-level
+		// resources, must like schemas.
+		if len(g) == 1 {
+			if r.Objects, err = excludeObjects(r.Objects, g); err != nil {
+				return nil, err
+			}
+		}
+	}
+	var schemas []*Schema
 Filter:
 	for _, s := range r.Schemas {
 		for i, g := range globs {
 			if len(g) > 3 {
 				return nil, fmt.Errorf("too many parts in pattern: %q", patterns[i])
 			}
-			match, err := filepath.Match(g[0], s.Name)
-			if err != nil {
-				return nil, err
-			}
-			if match {
-				// In case there is a match, and it is
-				// a single glob we exclude this
-				if len(g) == 1 {
-					continue Filter
-				}
-				if err := excludeS(s, g[1:]); err != nil {
+			if globS, exclude := excludeType(typeS, g[0]); exclude {
+				match, err := filepath.Match(globS, s.Name)
+				if err != nil {
 					return nil, err
 				}
+				if match {
+					// In case there is a match, and it is
+					// a single glob we exclude this
+					if len(g) == 1 {
+						continue Filter
+					}
+					if err := excludeS(s, g[1:]); err != nil {
+						return nil, err
+					}
+				}
 			}
 		}
 		schemas = append(schemas, s)
@@ -88,42 +99,10 @@ func split(patterns []string) ([][]string, error) {
 	return globs, nil
 }
 
-func excludeS(s *Schema, glob []string) error {
-	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, o)
-			}
-		} else {
-			objects = append(objects, o)
-		}
+func excludeS(s *Schema, glob []string) (err error) {
+	if s.Objects, err = excludeObjects(s.Objects, glob); err != nil {
+		return err
 	}
-	s.Objects = objects
 	if globT, exclude := excludeType(typeT, glob[0]); exclude {
 		var tables []*Table
 		for _, t := range s.Tables {
@@ -269,9 +248,48 @@ func excludeV(v *View, pattern string) (err error) {
 	return
 }
 
+func excludeObjects(all []Object, glob []string) ([]Object, error) {
+	var (
+		objects = make([]Object, 0, len(all))
+		t2glob  = make(map[string]struct {
+			glob    string
+			exclude bool
+		})
+	)
+	for _, o := range all {
+		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 nil, err
+			}
+			// No match or glob has more than one pattern.
+			if !match || len(glob) != 1 {
+				objects = append(objects, o)
+			}
+		} else {
+			objects = append(objects, o)
+		}
+	}
+	return objects, nil
+}
+
 const (
 	typeV  = "view"
 	typeT  = "table"
+	typeS  = "schema"
 	typeC  = "column"
 	typeI  = "index"
 	typeF  = "fk"
@@ -281,7 +299,7 @@ const (
 	typePr = "procedure"
 )
 
-var reType = regexp.MustCompile(`\[type=([a-z|]+)+\]$`)
+var reType = regexp.MustCompile(`\[type=([a-z|_]+)+\]$`)
 
 func excludeType(t, v string) (string, bool) {
 	matches := reType.FindStringSubmatch(v)


Test patch:
diff --git a/sql/schema/exclude_test.go b/sql/schema/exclude_test.go
index 7992b62e..29ab7ed4 100644
--- a/sql/schema/exclude_test.go
+++ b/sql/schema/exclude_test.go
@@ -456,3 +456,56 @@ func TestExcludeSchema(t *testing.T) {
 	require.Len(t, r.Schemas, 1)
 	require.Len(t, r.Schemas[0].Tables, 1)
 }
+
+// TestingUser is fake Realm objects for testing.
+type TestingUser struct {
+	Object
+	Name string
+}
+
+// SpecType returns the type of the user.
+func (t *TestingUser) SpecType() string {
+	return "testing_user"
+}
+
+// SpecName returns the name of the extension.
+func (t *TestingUser) SpecName() string {
+	return t.Name
+}
+
+func TestExcludeRealmObjects(t *testing.T) {
+	r := NewRealm(New("s1"), New("s2")).AddObjects(&TestingUser{Name: "s1"}, &TestingUser{Name: "s2"})
+	r, err := ExcludeRealm(r, []string{"s1[type=schema]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 1)
+	require.Equal(t, "s2", r.Schemas[0].Name)
+	require.Len(t, r.Objects, 2)
+
+	r, err = ExcludeRealm(r, []string{"s2[type=testing_user]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 1)
+	require.Len(t, r.Objects, 1)
+	require.Equal(t, "s1", r.Objects[0].(*TestingUser).Name)
+
+	r, err = ExcludeRealm(r, []string{"*"})
+	require.NoError(t, err)
+	require.Empty(t, r.Schemas)
+	require.Empty(t, r.Objects)
+
+	r = NewRealm(New("s1"), New("s2")).AddObjects(&TestingUser{Name: "s1"}, &TestingUser{Name: "s2"})
+	r, err = ExcludeRealm(r, []string{"*.*"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 2)
+	require.Len(t, r.Objects, 2)
+
+	r, err = ExcludeRealm(r, []string{"s*"})
+	require.NoError(t, err)
+	require.Empty(t, r.Schemas)
+	require.Empty(t, r.Objects)
+
+	r = NewRealm(New("s1"), New("s2")).AddObjects(&TestingUser{Name: "s1"}, &TestingUser{Name: "s2"})
+	r, err = ExcludeRealm(r, []string{"s*[type=testing_user]"})
+	require.NoError(t, err)
+	require.Len(t, r.Schemas, 2)
+	require.Empty(t, r.Objects)
+}