Add exclude schemas

Author: bplunkett-stripe

internal/migration_acceptance_tests/acceptance_test.go

@@ -130,7 +130,6 @@ func (suite *acceptanceTestSuite) runSubtest(tc acceptanceTestCase, expects expe
 			return diff.Generate(ctx, connPool, diff.DDLSchemaSource(newSchemaDDL),
 				append(planOpts,
 					diff.WithBetaDoNotCallWithAllowCustomSchemaOpts(),
-					diff.WithSchemas("public"),
 					diff.WithTempDbFactory(tempDbFactory),
 				)...)
 		}

internal/schema/filters.go

@@ -10,6 +10,12 @@ func schemaNameFilter(schema string) nameFilter {
 	}
 }
 
+func notSchemaNameFilter(schema string) nameFilter {
+	return func(obj SchemaQualifiedName) bool {
+		return obj.SchemaName != schema
+	}
+}
+
 func orNameFilter(filters ...nameFilter) nameFilter {
 	return func(obj SchemaQualifiedName) bool {
 		for _, filter := range filters {
@@ -21,6 +27,21 @@ func orNameFilter(filters ...nameFilter) nameFilter {
 	}
 }
 
+func andNameFilter(filters ...nameFilter) nameFilter {
+	return func(obj SchemaQualifiedName) bool {
+		if len(filters) == 0 {
+			return false
+		}
+
+		for _, filter := range filters {
+			if !filter(obj) {
+				return false
+			}
+		}
+		return true
+	}
+}
+
 func filterSliceByName[T any](objs []T, getNameFn func(T) SchemaQualifiedName, filter nameFilter) []T {
 	var filteredObjs []T
 	for _, obj := range objs {

internal/schema/filters_test.go

@@ -30,6 +30,70 @@ func (f fakeNameFilter) filter(input SchemaQualifiedName) bool {
 }
 
 func TestOrNameFilters(t *testing.T) {
+	someName1 := SchemaQualifiedName{
+		SchemaName:  "some_schema",
+		EscapedName: "some_name",
+	}
+	for _, tc := range []struct {
+		name        string
+		input       SchemaQualifiedName
+		filters     []fakeNameFilterMock
+		expectedOut bool
+	}{
+		{
+			name:        "empty",
+			input:       someName1,
+			expectedOut: false,
+		},
+		{
+			name:  "one filter (true)",
+			input: someName1,
+			filters: []fakeNameFilterMock{
+				{
+					expectedInput: someName1,
+					returnValue:   true,
+				},
+			},
+			expectedOut: true,
+		},
+		{
+			name:  "one filter (false)",
+			input: someName1,
+			filters: []fakeNameFilterMock{
+				{expectedInput: someName1, returnValue: false},
+			},
+			expectedOut: false,
+		},
+		{
+			name:  "two filters (false, true)",
+			input: someName1,
+			filters: []fakeNameFilterMock{
+				{expectedInput: someName1, returnValue: false},
+				{expectedInput: someName1, returnValue: true},
+			},
+			expectedOut: false,
+		},
+		{
+			name:  "two filters (true, true)",
+			input: someName1,
+			filters: []fakeNameFilterMock{
+				{expectedInput: someName1, returnValue: true},
+				{expectedInput: someName1, returnValue: true},
+			},
+			expectedOut: true,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var filters []nameFilter
+			for _, filter := range tc.filters {
+				filters = append(filters, newFakeNameFilter(t, filter).filter)
+			}
+			assert.Equal(t, tc.expectedOut, andNameFilter(filters...)(tc.input))
+		})
+	}
+}
+
+func TestAndNameFilters(t *testing.T) {
 	someName1 := SchemaQualifiedName{
 		SchemaName:  "some_schema",
 		EscapedName: "some_name",
@@ -91,5 +155,4 @@ func TestOrNameFilters(t *testing.T) {
 			assert.Equal(t, tc.expectedOut, orNameFilter(filters...)(tc.input))
 		})
 	}
-
 }

internal/schema/schema.go

@@ -155,7 +155,7 @@ func (t Table) IsPartitioned() bool {
 
 // IsPartition returns whether the table is a partition.
 // It represents a mismatch in modeling because the ForValues and ParentTable are stored separately.
-// Instead, they fields be stored under the same struct as a nilable pointer and this function should be deleted
+// Instead, the fields should be stored under the same struct as a nilable pointer, and this function should be deleted.
 func (t Table) IsPartition() bool {
 	return t.ParentTable != nil
 }
@@ -346,21 +346,31 @@ type (
 	GetSchemaOpt func(*getSchemaOptions)
 )
 
-// WithSchemas filters the schema to only include the given schemas. This unions with any schemas that are already included
-// via WithSchemas.
-func WithSchemas(schemas ...string) GetSchemaOpt {
+// WithIncludeSchemas filters the schema to only include the given schemas. This unions with any schemas that are already included
+// via WithIncludeSchemas. If empty, then all schemas are included.
+func WithIncludeSchemas(schemas ...string) GetSchemaOpt {
 	return func(o *getSchemaOptions) {
 		for _, schema := range schemas {
 			o.includeSchemas = append(o.includeSchemas, schema)
 		}
 	}
 }
 
+// WithExcludeSchemas filters the schema to exclude the given schemas. This unions with any schemas that are already excluded
+// via WithExcludeSchemas. If empty, then no schemas are excluded.
+func WithExcludeSchemas(schemas ...string) GetSchemaOpt {
+	return func(o *getSchemaOptions) {
+		o.excludeSchemas = append(o.excludeSchemas, schemas...)
+	}
+}
+
 type getSchemaOptions struct {
 	// includeSchemas is a list of schemas to include in the schema. If empty, then all schemas are included.
 	// We could have built a more complex set of options using the nameFilter system (nested unions and intersections);
 	// however, I felt it could expose some weird behaviors that we don't want to have to worry about just yet,
 	includeSchemas []string
+	// excludeSchemas is the exclude analog of includeSchemas.
+	excludeSchemas []string
 }
 
 // GetSchema fetches the database schema. It is a non-atomic operation.
@@ -382,34 +392,77 @@ func GetSchema(ctx context.Context, db queries.DBTX, opts ...GetSchemaOpt) (Sche
 		opt(&options)
 	}
 
+	nameFilter, err := buildNameFilter(options)
+	if err != nil {
+		return Schema{}, fmt.Errorf("building name filter: %w", err)
+	}
+
 	return (&schemaFetcher{
 		q:                      queries.New(db),
 		goroutineRunnerFactory: goroutineRunnerFactory,
-		nameFilter:             buildNameFilter(options),
+		nameFilter:             nameFilter,
 	}).getSchema(ctx)
 }
 
-func buildNameFilter(options getSchemaOptions) nameFilter {
-	if len(options.includeSchemas) == 0 {
+func buildNameFilter(options getSchemaOptions) (nameFilter, error) {
+	if intersection := intersect(options.includeSchemas, options.excludeSchemas); len(intersection) > 0 {
+		return nil, fmt.Errorf("schemas %v are both included and excluded", intersection)
+	}
+
+	includeSchemasFilter := buildIncludeSchemasFilter(options.includeSchemas)
+	excludeSchemasFilter := buildExcludeSchemasFilter(options.excludeSchemas)
+	return andNameFilter(includeSchemasFilter, excludeSchemasFilter), nil
+}
+
+func intersect(a, b []string) []string {
+	inAByA := make(map[string]bool)
+	for _, s := range a {
+		inAByA[s] = true
+	}
+	intersection := make([]string, 0, len(b))
+	for _, s := range b {
+		if inAByA[s] {
+			intersection = append(intersection, s)
+		}
+	}
+	return intersection
+}
+
+func buildIncludeSchemasFilter(schemas []string) nameFilter {
+	if len(schemas) == 0 {
 		return func(name SchemaQualifiedName) bool {
 			return true
 		}
 	}
 
 	var filters []nameFilter
-	for _, schema := range options.includeSchemas {
+	for _, schema := range schemas {
 		filters = append(filters, schemaNameFilter(schema))
 	}
 	return orNameFilter(filters...)
 }
 
+func buildExcludeSchemasFilter(schemas []string) nameFilter {
+	if len(schemas) == 0 {
+		return func(name SchemaQualifiedName) bool {
+			return true
+		}
+	}
+
+	var filters []nameFilter
+	for _, schema := range schemas {
+		filters = append(filters, notSchemaNameFilter(schema))
+	}
+	return andNameFilter(filters...)
+}
+
 type (
 	schemaFetcher struct {
 		q *queries.Queries
 		// goroutineRunnerFactory is a factory function that returns a GoroutineRunner. We need to be able to construct
 		// multiple GoroutineRunners to avoid deadlock created by circular dependencies of submitted go routines.
 		goroutineRunnerFactory func() concurrent.GoroutineRunner
-		// nameFilter is a filter that determienes which schema objects to include in the schema via their
+		// nameFilter is a filter that determines which schema objects to include in the schema via their
 		// schema name and object name.
 		//
 		// Currently, we don't do any sort of validation to ensure that all dependencies are included, so users might

internal/schema/schema_test.go

@@ -16,12 +16,13 @@ import (
 
 type testCase struct {
 	name           string
-	ddl            []string
 	opts           []GetSchemaOpt
+	ddl            []string
 	expectedSchema Schema
 	// expectedHash is the expected hash of the schema. If it is not provided, the test will not validate the hash.
-	expectedHash  string
-	expectedErrIs error
+	expectedHash        string
+	expectedErrIs       error
+	expectedErrContains string
 }
 
 var (
@@ -38,7 +39,7 @@ var (
 		{
 			name: "Simple schema (validate all schema objects and schema name filters)",
 			opts: []GetSchemaOpt{
-				WithSchemas("public", "schema_1", "schema_2"),
+				WithIncludeSchemas("public", "schema_1", "schema_2"),
 			},
 			ddl: []string{`
 			CREATE SCHEMA schema_1;
@@ -796,7 +797,7 @@ var (
 		},
 		{
 			name: "Filtering - filtering out the base table",
-			opts: []GetSchemaOpt{WithSchemas("public")},
+			opts: []GetSchemaOpt{WithIncludeSchemas("public")},
 			ddl: []string{`
 				CREATE SCHEMA schema_filtered_1;
 				CREATE TABLE schema_filtered_1.foobar(	
@@ -821,7 +822,7 @@ var (
 		},
 		{
 			name: "Filtering - filtering out partition",
-			opts: []GetSchemaOpt{WithSchemas("public")},
+			opts: []GetSchemaOpt{WithIncludeSchemas("public")},
 			ddl: []string{`
 				CREATE SCHEMA schema_filtered_1;
 				CREATE TABLE foobar(	
@@ -901,6 +902,69 @@ var (
 				},
 			},
 		},
+		{
+			name: "Filters - exclude schemas",
+			opts: []GetSchemaOpt{
+				WithExcludeSchemas("schema_1"),
+			},
+			ddl: []string{`
+				CREATE TABLE foobar();
+				CREATE SCHEMA schema_1;
+				CREATE TABLE schema_1.foobar();
+				CREATE SCHEMA schema_2;
+				CREATE TABLE schema_2.foobar();
+			`},
+			expectedSchema: Schema{
+				Tables: []Table{
+					{
+						SchemaQualifiedName: SchemaQualifiedName{SchemaName: "public", EscapedName: "\"foobar\""},
+						ReplicaIdentity:     ReplicaIdentityDefault,
+					},
+					{
+						SchemaQualifiedName: SchemaQualifiedName{SchemaName: "schema_2", EscapedName: "\"foobar\""},
+						ReplicaIdentity:     ReplicaIdentityDefault,
+					},
+				},
+			},
+		},
+		{
+			name: "Filters - include and exclude schemas",
+			opts: []GetSchemaOpt{
+				WithIncludeSchemas("schema_1"),
+				// schema_3 is inherently excluded since it is not included
+				WithExcludeSchemas("schema_2"),
+			},
+			ddl: []string{`
+				CREATE TABLE foobar();
+				CREATE SCHEMA schema_1;	
+				CREATE TABLE schema_1.foobar();
+				CREATE SCHEMA schema_2;
+				CREATE TABLE schema_2.foobar();
+			`},
+			expectedSchema: Schema{
+				Tables: []Table{
+					{
+						SchemaQualifiedName: SchemaQualifiedName{SchemaName: "schema_1", EscapedName: "\"foobar\""},
+						ReplicaIdentity:     ReplicaIdentityDefault,
+					},
+				},
+			},
+		},
+		{
+			name: "Filter - include and exclude a schema",
+			opts: []GetSchemaOpt{
+				WithIncludeSchemas("schema_1"),
+
+				WithIncludeSchemas("schema_2"),
+				WithExcludeSchemas("schema_2"),
+
+				WithExcludeSchemas("schema_3"),
+
+				WithExcludeSchemas("schema_4"),
+				WithExcludeSchemas("schema_4"),
+			},
+			expectedErrContains: "are both included and excluded",
+		},
 	}
 )
 
@@ -948,7 +1012,12 @@ func runTestCase(t *testing.T, engine *pgengine.Engine, testCase *testCase, getD
 	if testCase.expectedErrIs != nil {
 		require.ErrorIs(t, err, testCase.expectedErrIs)
 		return
-	} else {
+	}
+	if testCase.expectedErrContains != "" {
+		require.ErrorContains(t, err, testCase.expectedErrContains)
+		return
+	}
+	if testCase.expectedErrIs == nil && testCase.expectedErrContains == "" {
 		require.NoError(t, err)
 	}