From be692f8cc94d74033494d140c310e932c705e785 Mon Sep 17 00:00:00 2001
From: Mike Palmiotto <mike.palmiotto@crunchydata.com>
Date: Wed, 29 Mar 2017 14:59:37 +0000
Subject: [PATCH 2/2] Add partitioned table support to sepgsql

Account for RELKIND_PARTITIONED_RELATIONS in sepgsql and treat the objects like
regular relations. This allows for proper object_access hook behavior for
partitioned tables.
---
 contrib/sepgsql/label.c    |  3 ++-
 contrib/sepgsql/relation.c | 33 +++++++++++++++++++++++----------
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/contrib/sepgsql/label.c b/contrib/sepgsql/label.c
index 8a72503..c66581a 100644
--- a/contrib/sepgsql/label.c
+++ b/contrib/sepgsql/label.c
@@ -787,7 +787,8 @@ exec_object_restorecon(struct selabel_handle * sehnd, Oid catalogId)
 			case RelationRelationId:
 				relForm = (Form_pg_class) GETSTRUCT(tuple);
 
-				if (relForm->relkind == RELKIND_RELATION)
+				if (relForm->relkind == RELKIND_RELATION ||
+					relForm->relkind == RELKIND_PARTITIONED_TABLE)
 					objtype = SELABEL_DB_TABLE;
 				else if (relForm->relkind == RELKIND_SEQUENCE)
 					objtype = SELABEL_DB_SEQUENCE;
diff --git a/contrib/sepgsql/relation.c b/contrib/sepgsql/relation.c
index 4dc48a0..4fcebc1 100644
--- a/contrib/sepgsql/relation.c
+++ b/contrib/sepgsql/relation.c
@@ -54,12 +54,14 @@ sepgsql_attribute_post_create(Oid relOid, AttrNumber attnum)
 	ObjectAddress object;
 	Form_pg_attribute attForm;
 	StringInfoData audit_name;
+	char		relkind;
 
 	/*
 	 * Only attributes within regular relation have individual security
 	 * labels.
 	 */
-	if (get_rel_relkind(relOid) != RELKIND_RELATION)
+	relkind = get_rel_relkind(relOid);
+	if (relkind != RELKIND_RELATION && relkind != RELKIND_PARTITIONED_TABLE)
 		return;
 
 	/*
@@ -135,8 +137,10 @@ sepgsql_attribute_drop(Oid relOid, AttrNumber attnum)
 {
 	ObjectAddress object;
 	char	   *audit_name;
+	char		relkind;
 
-	if (get_rel_relkind(relOid) != RELKIND_RELATION)
+	relkind = get_rel_relkind(relOid);
+	if (relkind != RELKIND_RELATION && relkind != RELKIND_PARTITIONED_TABLE)
 		return;
 
 	/*
@@ -167,8 +171,11 @@ sepgsql_attribute_relabel(Oid relOid, AttrNumber attnum,
 {
 	ObjectAddress object;
 	char	   *audit_name;
+	char		relkind;
 
-	if (get_rel_relkind(relOid) != RELKIND_RELATION)
+	relkind = get_rel_relkind(relOid);
+
+	if (relkind != RELKIND_RELATION && relkind != RELKIND_PARTITIONED_TABLE)
 		ereport(ERROR,
 				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
 				 errmsg("cannot set security label on non-regular columns")));
@@ -209,8 +216,11 @@ sepgsql_attribute_setattr(Oid relOid, AttrNumber attnum)
 {
 	ObjectAddress object;
 	char	   *audit_name;
+	char		relkind;
+
+	relkind = get_rel_relkind(relOid);
 
-	if (get_rel_relkind(relOid) != RELKIND_RELATION)
+	if (relkind != RELKIND_RELATION && relkind != RELKIND_PARTITIONED_TABLE)
 		return;
 
 	/*
@@ -290,6 +300,7 @@ sepgsql_relation_post_create(Oid relOid)
 
 	switch (classForm->relkind)
 	{
+		case RELKIND_PARTITIONED_TABLE:
 		case RELKIND_RELATION:
 			tclass = SEPG_CLASS_DB_TABLE;
 			break;
@@ -336,7 +347,8 @@ sepgsql_relation_post_create(Oid relOid)
 								  true);
 
 	/*
-	 * Assign the default security label on the new relation
+	 * Assign the default security label on the new relation or partitioned
+	 * table.
 	 */
 	object.classId = RelationRelationId;
 	object.objectId = relOid;
@@ -344,10 +356,10 @@ sepgsql_relation_post_create(Oid relOid)
 	SetSecurityLabel(&object, SEPGSQL_LABEL_TAG, rcontext);
 
 	/*
-	 * We also assigns a default security label on columns of the new regular
-	 * tables.
+	 * We also assign a default security label on columns of a new table.
 	 */
-	if (classForm->relkind == RELKIND_RELATION)
+	if (classForm->relkind == RELKIND_RELATION ||
+		classForm->relkind == RELKIND_PARTITIONED_TABLE)
 	{
 		Relation	arel;
 		ScanKeyData akey;
@@ -422,6 +434,7 @@ sepgsql_relation_drop(Oid relOid)
 	relkind = get_rel_relkind(relOid);
 	switch (relkind)
 	{
+		case RELKIND_PARTITIONED_TABLE:
 		case RELKIND_RELATION:
 			tclass = SEPG_CLASS_DB_TABLE;
 			break;
@@ -485,7 +498,7 @@ sepgsql_relation_drop(Oid relOid)
 	/*
 	 * check db_column:{drop} permission
 	 */
-	if (relkind == RELKIND_RELATION)
+	if (relkind == RELKIND_RELATION || relkind == RELKIND_PARTITIONED_TABLE)
 	{
 		Form_pg_attribute attForm;
 		CatCList   *attrList;
@@ -531,7 +544,7 @@ sepgsql_relation_relabel(Oid relOid, const char *seclabel)
 	uint16_t	tclass = 0;
 
 	relkind = get_rel_relkind(relOid);
-	if (relkind == RELKIND_RELATION)
+	if (relkind == RELKIND_RELATION || relkind == RELKIND_PARTITIONED_TABLE)
 		tclass = SEPG_CLASS_DB_TABLE;
 	else if (relkind == RELKIND_SEQUENCE)
 		tclass = SEPG_CLASS_DB_SEQUENCE;
-- 
1.8.3.1