Support force row level security on table owner

In GPDB 7+, users can add row level security to their tables to enforce row permissions for other users. On top of that, the table owning user can force the restriction on themselves. Authored-by: Kevin Yeap <kyeap@vmware.com> Co-authored-by: Andrew Repp <reppa@vmware.com>
silvertl · Nov 22, 2022 · 49df00b · 49df00b
1 parent 526b859
commit 49df00b
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 6 deletions.
diff --git a/backup/predata_relations.go b/backup/predata_relations.go
@@ -151,7 +151,7 @@ func PrintRegularTableCreateStatement(metadataFile *utils.FileWithByteCount, toc
 	if table.PartitionKeyDef != "" {
 		metadataFile.MustPrintf("PARTITION BY %s ", table.PartitionKeyDef)
 	}
-	if len(table.Inherits) != 0  && table.AttachPartitionInfo == (AttachPartitionInfo{}) {
+	if len(table.Inherits) != 0 && table.AttachPartitionInfo == (AttachPartitionInfo{}) {
 		dependencyList := strings.Join(table.Inherits, ", ")
 		metadataFile.MustPrintf("INHERITS (%s) ", dependencyList)
 	}
@@ -199,7 +199,7 @@ func printColumnDefinitions(metadataFile *utils.FileWithByteCount, columnDefs []
 			line += fmt.Sprintf(" COLLATE %s", column.Collation)
 		}
 		if column.HasDefault {
-			if column.AttGenerated != ""  {
+			if column.AttGenerated != "" {
 				line += fmt.Sprintf(" GENERATED ALWAYS AS %s %s", column.DefaultVal, column.AttGenerated)
 			} else {
 				line += fmt.Sprintf(" DEFAULT %s", column.DefaultVal)
@@ -274,10 +274,16 @@ func PrintPostCreateTableStatements(metadataFile *utils.FileWithByteCount, toc *
 				utils.MakeFQN(alteredPartitionRelation.OldSchema, alteredPartitionRelation.Name), alteredPartitionRelation.NewSchema))
 	}
 
-	attachInfo := table.AttachPartitionInfo
-	if (attachInfo != AttachPartitionInfo{}) {
-		statements = append(statements,
-			fmt.Sprintf("ALTER TABLE ONLY %s ATTACH PARTITION %s %s;", table.Inherits[0], attachInfo.Relname, attachInfo.Expr))
+	if connectionPool.Version.AtLeast("7") {
+		attachInfo := table.AttachPartitionInfo
+		if (attachInfo != AttachPartitionInfo{}) {
+			statements = append(statements,
+				fmt.Sprintf("ALTER TABLE ONLY %s ATTACH PARTITION %s %s;", table.Inherits[0], attachInfo.Relname, attachInfo.Expr))
+		}
+
+		if table.ForceRowSecurity {
+			statements = append(statements, fmt.Sprintf("ALTER TABLE ONLY %s FORCE ROW LEVEL SECURITY;", table.FQN()))
+		}
 	}
 
 	PrintStatements(metadataFile, toc, table, statements)

diff --git a/backup/predata_relations_tables_test.go b/backup/predata_relations_tables_test.go
@@ -688,5 +688,11 @@ ALTER TABLE schema1.table1 SET SCHEMA schema2;
 
 ALTER TABLE schema2.table2 SET SCHEMA schema1;`)
 		})
+		It("prints force row security", func() {
+			testutils.SkipIfBefore7(connectionPool)
+			testTable.ForceRowSecurity = true
+			backup.PrintPostCreateTableStatements(backupfile, tocfile, testTable, backup.ObjectMetadata{})
+			testhelper.ExpectRegexp(buffer, `ALTER TABLE ONLY public.tablename FORCE ROW LEVEL SECURITY;`)
+		})
 	})
 })
diff --git a/backup/queries_table_defs.go b/backup/queries_table_defs.go
@@ -80,6 +80,7 @@ type TableDefinition struct {
 	AccessMethodName        string
 	PartitionKeyDef         string
 	AttachPartitionInfo     AttachPartitionInfo
+	ForceRowSecurity        bool
 }
 
 /*
@@ -106,6 +107,7 @@ func ConstructDefinitionsForTables(connectionPool *dbconn.DBConn, tableRelations
 	partitionAlteredSchemaMap := GetPartitionAlteredSchema(connectionPool)
 	partitionKeyDefs := GetPartitionKeyDefs(connectionPool)
 	attachPartitionInfo := GetAttachPartitionInfo(connectionPool)
+	forceRowSecurity := GetForceRowSecurity(connectionPool)
 
 	gplog.Verbose("Constructing table definition map")
 	for _, tableRel := range tableRelations {
@@ -129,6 +131,7 @@ func ConstructDefinitionsForTables(connectionPool *dbconn.DBConn, tableRelations
 			AccessMethodName:        accessMethodMap[oid],
 			PartitionKeyDef:         partitionKeyDefs[oid],
 			AttachPartitionInfo:     attachPartitionInfo[oid],
+			ForceRowSecurity:        forceRowSecurity[oid],
 		}
 		if tableDef.Inherits == nil {
 			tableDef.Inherits = []string{}
@@ -638,6 +641,28 @@ func GetAttachPartitionInfo(connectionPool *dbconn.DBConn) map[uint32]AttachPart
 	return resultMap
 }
 
+func GetForceRowSecurity(connectionPool *dbconn.DBConn) map[uint32]bool {
+	resultMap := make(map[uint32]bool)
+	if connectionPool.Version.Before("7") {
+		return resultMap
+	}
+
+	query := fmt.Sprintf(`
+	SELECT oid
+	FROM pg_class
+	WHERE relforcerowsecurity = 't'
+		AND oid >= %d`, FIRST_NORMAL_OBJECT_ID)
+
+	results := make([]uint32, 0)
+	err := connectionPool.Select(&results, query)
+	gplog.FatalOnError(err)
+
+	for _, oid := range results {
+		resultMap[oid] = true
+	}
+	return resultMap
+}
+
 func selectAsOidToStringMap(connectionPool *dbconn.DBConn, query string) map[uint32]string {
 	var results []struct {
 		Oid   uint32

diff --git a/integration/predata_relations_create_test.go b/integration/predata_relations_create_test.go
@@ -462,6 +462,16 @@ SET SUBPARTITION TEMPLATE ` + `
 			testChildTable.AttachPartitionInfo.Oid = childTableOid
 			structmatcher.ExpectStructsToMatch(&testChildTable.AttachPartitionInfo, attachPartitionInfoMap[childTableOid])
 		})
+		It("prints an ALTER statement to force row level security on the table owner", func() {
+			testutils.SkipIfBefore7(connectionPool)
+
+			testTable.ForceRowSecurity = true
+			backup.PrintPostCreateTableStatements(backupfile, tocfile, testTable, tableMetadata)
+			testhelper.AssertQueryRuns(connectionPool, buffer.String())
+			testTable.Oid = testutils.OidFromObjectName(connectionPool, "public", "testtable", backup.TYPE_RELATION)
+			resultTable := backup.ConstructDefinitionsForTables(connectionPool, []backup.Relation{testTable.Relation})[0]
+			Expect(resultTable.ForceRowSecurity).To(Equal(true))
+		})
 	})
 	Describe("PrintCreateViewStatements", func() {
 		var viewDef sql.NullString