fixup! PG-1504 Make partitions inherit encryption status

Author: Anders Åstrand

contrib/pg_tde/expected/partition_table.out

@@ -86,7 +86,8 @@ SELECT pg_tde_is_encrypted('partition_q4_2024');
  t
 (1 row)
 
--- Partition inherits encryption status from parent table
+-- Partition inherits encryption status from parent table if default is heap and parent is tde_heap
+SET default_table_access_method = "heap";
 CREATE TABLE partition_inherit_parent (a int) PARTITION BY RANGE (a) USING tde_heap;
 CREATE TABLE partition_inherit_child PARTITION OF partition_inherit_parent FOR VALUES FROM (0) TO (10);
 SELECT pg_tde_is_encrypted('partition_inherit_child');
@@ -96,5 +97,18 @@ SELECT pg_tde_is_encrypted('partition_inherit_child');
 (1 row)
 
 DROP TABLE partition_inherit_parent;
+RESET default_table_access_method;
+-- Partition inherits encryption status from parent table if default is tde_heap and parent is heap
+SET default_table_access_method = "tde_heap";
+CREATE TABLE partition_inherit_parent (a int) PARTITION BY RANGE (a) USING heap;
+CREATE TABLE partition_inherit_child PARTITION OF partition_inherit_parent FOR VALUES FROM (0) TO (10);
+SELECT pg_tde_is_encrypted('partition_inherit_child');
+ pg_tde_is_encrypted 
+---------------------
+ f
+(1 row)
+
+DROP TABLE partition_inherit_parent;
+RESET default_table_access_method;
 DROP TABLE partitioned_table;
 DROP EXTENSION pg_tde;

contrib/pg_tde/sql/partition_table.sql

@@ -31,11 +31,21 @@ SELECT pg_tde_is_encrypted('partition_q2_2024');
 SELECT pg_tde_is_encrypted('partition_q3_2024');
 SELECT pg_tde_is_encrypted('partition_q4_2024');
 
--- Partition inherits encryption status from parent table
+-- Partition inherits encryption status from parent table if default is heap and parent is tde_heap
+SET default_table_access_method = "heap";
 CREATE TABLE partition_inherit_parent (a int) PARTITION BY RANGE (a) USING tde_heap;
 CREATE TABLE partition_inherit_child PARTITION OF partition_inherit_parent FOR VALUES FROM (0) TO (10);
 SELECT pg_tde_is_encrypted('partition_inherit_child');
 DROP TABLE partition_inherit_parent;
+RESET default_table_access_method;
+
+-- Partition inherits encryption status from parent table if default is tde_heap and parent is heap
+SET default_table_access_method = "tde_heap";
+CREATE TABLE partition_inherit_parent (a int) PARTITION BY RANGE (a) USING heap;
+CREATE TABLE partition_inherit_child PARTITION OF partition_inherit_parent FOR VALUES FROM (0) TO (10);
+SELECT pg_tde_is_encrypted('partition_inherit_child');
+DROP TABLE partition_inherit_parent;
+RESET default_table_access_method;
 
 DROP TABLE partitioned_table;
 DROP EXTENSION pg_tde;

contrib/pg_tde/src/pg_tde_event_capture.c

@@ -155,16 +155,17 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 		validateCurrentEventTriggerState(true);
 		tdeCurrentCreateEvent.tid = GetCurrentFullTransactionId();
 
-		if (shouldEncryptTable(stmt->accessMethod))
+
+		if (stmt->accessMethod && strcmp(stmt->accessMethod, "tde_heap") == 0)
 		{
+			/* If access method is explicitly set to tde_heap, always encrypt. */
 			tdeCurrentCreateEvent.encryptMode = true;
 		}
 		else if (!stmt->accessMethod && stmt->partbound)
 		{
-
 			/*
-			 * If this is a partition of a parent table, and no access method
-			 * is specified, access method will be inherited and we need to
+			 * If no access method is specified, and this is a partition of a
+			 * parent table, access method will be inherited and we need to
 			 * deal with setting the encryption status properly.
 			 *
 			 * AccessExclusiveLock might seem excessive, but it's what
@@ -184,6 +185,15 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 				tdeCurrentCreateEvent.encryptMode = true;
 			}
 		}
+		else if (!stmt->accessMethod && strcmp(default_table_access_method, "tde_heap") == 0)
+		{
+			/*
+			 * If no access method is specified, and this is not a partition
+			 * of a parent table, refer to the default access method to set
+			 * encryption status.
+			 */
+			tdeCurrentCreateEvent.encryptMode = true;
+		}
 
 		checkEncryptionStatus();
 	}