Add RLS policy support (#132)

stripe · May 29, 2024 · 045fd2c · 045fd2c
1 parent 0ae5788
commit 045fd2c
Show file tree

Hide file tree

Showing 28 changed files with 1,822 additions and 108 deletions.
diff --git a/internal/migration_acceptance_tests/acceptance_test.go b/internal/migration_acceptance_tests/acceptance_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"reflect"
 	"testing"
 
 	"github.com/google/uuid"
@@ -42,7 +43,9 @@ type (
 	planFactory func(ctx context.Context, connPool sqldb.Queryable, tempDbFactory tempdb.Factory, newSchemaDDL []string, opts ...diff.PlanOpt) (diff.Plan, error)
 
 	acceptanceTestCase struct {
-		name         string
+		name string
+		// roles is a list of roles that should be created before the DDL is applied
+		roles        []string
 		oldSchemaDDL []string
 		newSchemaDDL []string
 
@@ -88,12 +91,16 @@ func (suite *acceptanceTestSuite) runTestCases(acceptanceTestCases []acceptanceT
 			suite.Run("vanilla", func() {
 				suite.runSubtest(tc, tc.vanillaExpectations, nil)
 			})
-			suite.Run("with data packing (and ignoring column order)", func() {
-				suite.runSubtest(tc, tc.dataPackingExpectations, []diff.PlanOpt{
-					diff.WithDataPackNewTables(),
-					diff.WithLogger(log.SimpleLogger()),
+			// Only run the data packing test if there are expectations for it. We should strip out the embedded
+			// data packing tests and make them their own tests to simplify the test suite.
+			if !reflect.DeepEqual(tc.dataPackingExpectations, expectations{}) {
+				suite.Run("with data packing (and ignoring column order)", func() {
+					suite.runSubtest(tc, tc.dataPackingExpectations, []diff.PlanOpt{
+						diff.WithDataPackNewTables(),
+						diff.WithLogger(log.SimpleLogger()),
+					})
 				})
-			})
+			}
 		})
 	}
 }
@@ -106,6 +113,19 @@ func (suite *acceptanceTestSuite) runSubtest(tc acceptanceTestCase, expects expe
 		expects.outputState = tc.newSchemaDDL
 	}
 
+	// Create roles since they are global
+	rootDb, err := sql.Open("pgx", suite.pgEngine.GetPostgresDatabaseDSN())
+	suite.Require().NoError(err)
+	defer rootDb.Close()
+	for _, r := range tc.roles {
+		_, err := rootDb.Exec(fmt.Sprintf("CREATE ROLE %s", r))
+		suite.Require().NoError(err)
+	}
+	defer func() {
+		// This will drop the roles (and attempt to reset other cluster-level state)
+		suite.Require().NoError(pgengine.ResetInstance(context.Background(), rootDb))
+	}()
+
 	// Apply old schema DDL to old DB
 	oldDb, err := suite.pgEngine.CreateDatabase()
 	suite.Require().NoError(err)

diff --git a/internal/migration_acceptance_tests/backwards_compat_cases_test.go b/internal/migration_acceptance_tests/backwards_compat_cases_test.go
@@ -177,6 +177,6 @@ var backCompatAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestBackCompatAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestBackCompatTestCases() {
 	suite.runTestCases(backCompatAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/check_constraint_cases_test.go b/internal/migration_acceptance_tests/check_constraint_cases_test.go
@@ -568,6 +568,6 @@ var checkConstraintCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestCheckConstraintAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestCheckConstraintTestCases() {
 	suite.runTestCases(checkConstraintCases)
 }
diff --git a/internal/migration_acceptance_tests/column_cases_test.go b/internal/migration_acceptance_tests/column_cases_test.go
@@ -945,6 +945,6 @@ var columnAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestColumnAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestColumnTestCases() {
 	suite.runTestCases(columnAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/enum_cases_test.go b/internal/migration_acceptance_tests/enum_cases_test.go
@@ -126,6 +126,6 @@ var enumAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (s *acceptanceTestSuite) TestEnumTestCases() {
-	s.runTestCases(enumAcceptanceTestCases)
+func (suite *acceptanceTestSuite) TestEnumTestCases() {
+	suite.runTestCases(enumAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/extensions_cases_test.go b/internal/migration_acceptance_tests/extensions_cases_test.go
@@ -72,6 +72,6 @@ var extensionAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestExtensionAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestExtensionTestCases() {
 	suite.runTestCases(extensionAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/function_cases_test.go b/internal/migration_acceptance_tests/function_cases_test.go
@@ -577,6 +577,6 @@ var functionAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestFunctionAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestFunctionTestCases() {
 	suite.runTestCases(functionAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/index_cases_test.go b/internal/migration_acceptance_tests/index_cases_test.go
@@ -803,6 +803,6 @@ var indexAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestIndexAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestIndexTestCases() {
 	suite.runTestCases(indexAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/local_partition_index_cases_test.go b/internal/migration_acceptance_tests/local_partition_index_cases_test.go
@@ -605,6 +605,6 @@ var localPartitionIndexAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestLocalPartitionIndexAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestLocalPartitionIndexTestCases() {
 	suite.runTestCases(localPartitionIndexAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/named_schema_cases_test.go b/internal/migration_acceptance_tests/named_schema_cases_test.go
@@ -40,6 +40,6 @@ var namedSchemaAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestNamedSchemaAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestNamedSchemaTestCases() {
 	suite.runTestCases(namedSchemaAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/partitioned_index_cases_test.go b/internal/migration_acceptance_tests/partitioned_index_cases_test.go
@@ -1245,6 +1245,6 @@ var partitionedIndexAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestPartitionedIndexAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestPartitionedIndexTestCases() {
 	suite.runTestCases(partitionedIndexAcceptanceTestCases)
 }
diff --git a/internal/migration_acceptance_tests/partitioned_table_cases_test.go b/internal/migration_acceptance_tests/partitioned_table_cases_test.go
@@ -19,9 +19,13 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 				UNIQUE (foo, bar)
 			) PARTITION BY LIST (foo);
 			ALTER TABLE foobar REPLICA IDENTITY FULL;
+			ALTER TABLE foobar ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE foobar FORCE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_1 PARTITION OF foobar FOR VALUES IN ('foo_1');
 			ALTER TABLE foobar_1 REPLICA IDENTITY DEFAULT ;
+			ALTER TABLE foobar_1 ENABLE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_2 PARTITION OF foobar FOR VALUES IN ('foo_2');
+			ALTER TABLE foobar_2 FORCE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_3 PARTITION OF foobar FOR VALUES IN ('foo_3');
 			-- partitioned indexes
 			CREATE UNIQUE INDEX foobar_unique_idx ON foobar(foo, bar);
@@ -59,10 +63,17 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 				UNIQUE (foo, bar)
 			) PARTITION BY LIST (foo);
 			ALTER TABLE foobar REPLICA IDENTITY FULL;
+			ALTER TABLE foobar ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE foobar FORCE ROW LEVEL SECURITY;
 			-- partitions
+			ALTER TABLE foobar REPLICA IDENTITY FULL;
+			ALTER TABLE foobar ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE foobar FORCE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_1 PARTITION OF foobar FOR VALUES IN ('foo_1');
 			ALTER TABLE foobar_1 REPLICA IDENTITY DEFAULT ;
+			ALTER TABLE foobar_1 ENABLE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_2 PARTITION OF foobar FOR VALUES IN ('foo_2');
+			ALTER TABLE foobar_2 FORCE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_3 PARTITION OF foobar FOR VALUES IN ('foo_3');
 			-- partitioned indexes
 			CREATE UNIQUE INDEX foobar_unique_idx ON foobar(foo, bar);
@@ -98,7 +109,7 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 		},
 	},
 	{
-		name:         "Create partitioned table with shared primary key",
+		name:         "Create partitioned table with shared primary key and RLS enabled globally",
 		oldSchemaDDL: nil,
 		newSchemaDDL: []string{
 			`
@@ -113,6 +124,9 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 			    UNIQUE (foo, bar)
 			) PARTITION BY LIST (foo);
 			ALTER TABLE schema_1."Foobar" REPLICA IDENTITY FULL;
+			
+			ALTER TABLE schema_1."Foobar" ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE schema_1."Foobar" FORCE ROW LEVEL SECURITY;
 
 			-- partitions
 			CREATE SCHEMA schema_2;
@@ -145,6 +159,9 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 			) PARTITION BY LIST (foo);
 			ALTER TABLE schema_1."Foobar" REPLICA IDENTITY FULL;
 
+			ALTER TABLE schema_1."Foobar" ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE schema_1."Foobar" FORCE ROW LEVEL SECURITY;
+
 			-- partitions
 			CREATE SCHEMA schema_2;
 			CREATE TABLE schema_2."FOOBAR_1" PARTITION OF schema_1."Foobar"(
@@ -165,7 +182,7 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 		},
 	},
 	{
-		name:         "Create partitioned table with local primary keys",
+		name:         "Create partitioned table with local primary keys and RLS enabled locally",
 		oldSchemaDDL: nil,
 		newSchemaDDL: []string{
 			`
@@ -182,9 +199,11 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 			    bar NOT NULL,
 			    PRIMARY KEY (foo, id)
 			) FOR VALUES IN ('foo_1');
+			ALTER TABLE "FOOBAR_1" ENABLE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_2 PARTITION OF "Foobar"(
 			    PRIMARY KEY (foo, bar)
 			) FOR VALUES IN ('foo_2');
+			ALTER TABLE foobar_2 FORCE ROW LEVEL SECURITY;
 			CREATE TABLE foobar_3 PARTITION OF "Foobar"(
 			    PRIMARY KEY (foo, fizz),
 			    UNIQUE (foo, bar)
@@ -334,6 +353,80 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 			diff.MigrationHazardTypeCorrectness,
 		},
 	},
+	{
+		name: "Enable RLS of parent and children",
+		oldSchemaDDL: []string{
+			`
+			CREATE TABLE foobar(
+			    id INT,
+				foo VARCHAR(255),
+			    PRIMARY KEY (foo, id)
+			) PARTITION BY LIST (foo);
+
+			-- partitions
+			CREATE TABLE foobar_1 PARTITION OF foobar FOR VALUES IN ('foo_1');
+			CREATE TABLE foobar_2 PARTITION OF foobar FOR VALUES IN ('foo_2');
+			CREATE TABLE foobar_3 PARTITION OF foobar FOR VALUES IN ('foo_3');
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+			CREATE TABLE foobar(
+			    id INT,
+				foo VARCHAR(255),
+			    PRIMARY KEY (foo, id)
+			) PARTITION BY LIST (foo);
+			ALTER TABLE foobar ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE foobar FORCE ROW LEVEL SECURITY;
+			-- partitions
+			CREATE TABLE foobar_1 PARTITION OF foobar FOR VALUES IN ('foo_1');
+			ALTER TABLE foobar_1 ENABLE ROW LEVEL SECURITY;
+			CREATE TABLE foobar_2 PARTITION OF foobar FOR VALUES IN ('foo_2');
+			ALTER TABLE foobar_2 FORCE ROW LEVEL SECURITY;
+			CREATE TABLE foobar_3 PARTITION OF foobar FOR VALUES IN ('foo_3');
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeAuthzUpdate,
+		},
+	},
+	{
+		name: "Disable RLS of parent and children",
+		oldSchemaDDL: []string{
+			`
+			CREATE TABLE foobar(
+			    id INT,
+				foo VARCHAR(255),
+			    PRIMARY KEY (foo, id)
+			) PARTITION BY LIST (foo);
+			ALTER TABLE foobar ENABLE ROW LEVEL SECURITY;
+			ALTER TABLE foobar FORCE ROW LEVEL SECURITY;
+			-- partitions
+			CREATE TABLE foobar_1 PARTITION OF foobar FOR VALUES IN ('foo_1');
+			ALTER TABLE foobar_1 ENABLE ROW LEVEL SECURITY;
+			CREATE TABLE foobar_2 PARTITION OF foobar FOR VALUES IN ('foo_2');
+			ALTER TABLE foobar_2 FORCE ROW LEVEL SECURITY;
+			CREATE TABLE foobar_3 PARTITION OF foobar FOR VALUES IN ('foo_3');
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+			CREATE TABLE foobar(
+			    id INT,
+				foo VARCHAR(255),
+			    PRIMARY KEY (foo, id)
+			) PARTITION BY LIST (foo);
+
+			-- partitions
+			CREATE TABLE foobar_1 PARTITION OF foobar FOR VALUES IN ('foo_1');
+			CREATE TABLE foobar_2 PARTITION OF foobar FOR VALUES IN ('foo_2');
+			CREATE TABLE foobar_3 PARTITION OF foobar FOR VALUES IN ('foo_3');
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeAuthzUpdate,
+		},
+	},
 	{
 		name: "Alter table: New primary key, new unique constraint, dropped unique constraint, change column types, delete partitioned index, new partitioned index, delete local index, add local index, validate check constraint, validate FK, delete FK",
 		oldSchemaDDL: []string{
@@ -1248,6 +1341,6 @@ var partitionedTableAcceptanceTestCases = []acceptanceTestCase{
 	},
 }
 
-func (suite *acceptanceTestSuite) TestPartitionedTableAcceptanceTestCases() {
+func (suite *acceptanceTestSuite) TestPartitionedTableTestCases() {
 	suite.runTestCases(partitionedTableAcceptanceTestCases)
 }