PG-1663 Make sure indexes on paritioned tables are encrypted

Since we only looked at the parent table and not on the whole tree when
setting the status of the encrypted indexes we could easily accidentally
create a plain text index on an encrypted table.

This patch also makes sure to disallow adding indexes to an inheritance
tree where the tables are a mix of encrypted and unecrypted tables.

contrib/pg_tde/expected/partition_table.out

@@ -161,4 +161,50 @@ SELECT pg_tde_is_encrypted('partition_child_tde_heap');
 
 DROP TABLE partition_parent;
 RESET pg_tde.enforce_encryption;
+-- Partitioned indexes should be encrypted
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE TABLE partition_child PARTITION OF partition_parent FOR VALUES FROM (0) TO (9) USING tde_heap;
+CREATE INDEX ON partition_parent (a);
+SELECT pg_tde_is_encrypted('partition_parent_a_idx'); -- Also check that the parent index is NULL
+ pg_tde_is_encrypted 
+---------------------
+ 
+(1 row)
+
+SELECT pg_tde_is_encrypted('partition_child_a_idx');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+DROP TABLE partition_parent;
+-- Partitioned indexes should be not encrypted with heap
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE TABLE partition_child PARTITION OF partition_parent FOR VALUES FROM (0) TO (9) USING heap;
+CREATE INDEX ON partition_parent (a);
+SELECT pg_tde_is_encrypted('partition_child_a_idx');
+ pg_tde_is_encrypted 
+---------------------
+ f
+(1 row)
+
+DROP TABLE partition_parent;
+-- We refuse to create an index when the inheritance heirarchy has mixed statuses
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE TABLE partition_child_heap PARTITION OF partition_parent FOR VALUES FROM (0) TO (9) USING heap;
+CREATE TABLE partition_child_tde_heap PARTITION OF partition_parent FOR VALUES FROM (10) TO (19) USING tde_heap;
+CREATE INDEX ON partition_parent (a);
+ERROR:  Recursive CREATE INDEX on a mix of encrypted and unencrypted relations is not supported
+DROP TABLE partition_parent;
+-- Index should also be encrypted for new partitionins
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE INDEX ON partition_parent (a);
+CREATE TABLE partition_child PARTITION OF partition_parent FOR VALUES FROM (10) TO (19) USING tde_heap;
+SELECT pg_tde_is_encrypted('partition_child_a_idx');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+DROP TABLE partition_parent;
 DROP EXTENSION pg_tde;

contrib/pg_tde/sql/partition_table.sql

@@ -80,4 +80,33 @@ SELECT pg_tde_is_encrypted('partition_child_tde_heap');
 DROP TABLE partition_parent;
 RESET pg_tde.enforce_encryption;
 
+-- Partitioned indexes should be encrypted
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE TABLE partition_child PARTITION OF partition_parent FOR VALUES FROM (0) TO (9) USING tde_heap;
+CREATE INDEX ON partition_parent (a);
+SELECT pg_tde_is_encrypted('partition_parent_a_idx'); -- Also check that the parent index is NULL
+SELECT pg_tde_is_encrypted('partition_child_a_idx');
+DROP TABLE partition_parent;
+
+-- Partitioned indexes should be not encrypted with heap
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE TABLE partition_child PARTITION OF partition_parent FOR VALUES FROM (0) TO (9) USING heap;
+CREATE INDEX ON partition_parent (a);
+SELECT pg_tde_is_encrypted('partition_child_a_idx');
+DROP TABLE partition_parent;
+
+-- We refuse to create an index when the inheritance heirarchy has mixed statuses
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE TABLE partition_child_heap PARTITION OF partition_parent FOR VALUES FROM (0) TO (9) USING heap;
+CREATE TABLE partition_child_tde_heap PARTITION OF partition_parent FOR VALUES FROM (10) TO (19) USING tde_heap;
+CREATE INDEX ON partition_parent (a);
+DROP TABLE partition_parent;
+
+-- Index should also be encrypted for new partitionins
+CREATE TABLE partition_parent (a int) PARTITION BY RANGE (a);
+CREATE INDEX ON partition_parent (a);
+CREATE TABLE partition_child PARTITION OF partition_parent FOR VALUES FROM (10) TO (19) USING tde_heap;
+SELECT pg_tde_is_encrypted('partition_child_a_idx');
+DROP TABLE partition_parent;
+
 DROP EXTENSION pg_tde;

contrib/pg_tde/src/pg_tde_event_capture.c

@@ -258,21 +258,22 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 	if (IsA(parsetree, IndexStmt))
 	{
 		IndexStmt  *stmt = castNode(IndexStmt, parsetree);
-		Relation	rel;
 		TdeDdlEvent event = {.parsetree = parsetree};
+		EncryptionMix encmix;
+		Oid			relid = RangeVarGetRelid(stmt->relation, AccessShareLock, false);
 
-		rel = table_openrv(stmt->relation, AccessShareLock);
+		encmix = alter_table_encryption_mix(relid);
 
-		if (rel->rd_rel->relam == get_tde_table_am_oid())
-		{
+		if (encmix == ENC_MIX_ENCRYPTED)
 			event.encryptMode = TDE_ENCRYPT_MODE_ENCRYPT;
-			checkPrincipalKeyConfigured();
-		}
-		else
+		else if (encmix == ENC_MIX_PLAIN)
 			event.encryptMode = TDE_ENCRYPT_MODE_PLAIN;
-
-		/* Hold on to lock until end of transaction */
-		table_close(rel, NoLock);
+		else if (encmix == ENC_MIX_UNKNOWN)
+			event.encryptMode = TDE_ENCRYPT_MODE_RETAIN;
+		else
+			ereport(ERROR,
+					errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					errmsg("Recursive CREATE INDEX on a mix of encrypted and unencrypted relations is not supported"));
 
 		push_event_stack(&event);
 	}