PG-1507 PG-1508 Disallow ALTER TABLE when there is a mix of encrypted and unencrypted

Since ALTER TABLE can modify multiple due to inheritance or typed tables
which can for example result in TOAST tables being created for some or all
of the modified tables while the event trigger is only fired once we cannot
rely on the event stack to make sure we get the correct encryption status.

Our solution is to be cautious and only modify tables when all tables
with storage are either encrypt or not encrypted. If there is a mix we
will throw an error. The result of this is also used to properly inform
the SMGR of the current encryption status. We apply the same logic to
ALTER TYPE for typed tables.

A possibility for future improvement would be to limit this check only
to commands which recurse down to child tables and/or can create new
relfilenodes.

contrib/pg_tde/expected/partition_table.out

@@ -52,6 +52,7 @@ SELECT pg_tde_is_encrypted('partition_q4_2024');
 (1 row)
 
 ALTER TABLE partitioned_table SET ACCESS METHOD heap;
+ERROR:  Recursive ALTER TABLE on a mix of encrypted and unencrypted relations is not supported
 ALTER TABLE partition_q1_2024 SET ACCESS METHOD heap;
 ALTER TABLE partition_q2_2024 SET ACCESS METHOD tde_heap;
 ALTER TABLE partition_q3_2024 SET ACCESS METHOD heap;
@@ -86,6 +87,12 @@ SELECT pg_tde_is_encrypted('partition_q4_2024');
  t
 (1 row)
 
+-- Does not care about parent AM as long as all children with storage use the same
+ALTER TABLE partition_q1_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partition_q2_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partition_q3_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partition_q4_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partitioned_table SET ACCESS METHOD heap;
 DROP TABLE partitioned_table;
 -- Partition inherits encryption status from parent table if default is heap and parent is tde_heap
 SET default_table_access_method = "heap";

contrib/pg_tde/expected/recreate_storage.out

@@ -229,6 +229,49 @@ SELECT pg_tde_is_encrypted('plain_table_id2_seq2');
  f
 (1 row)
 
+-- Enforce that we do not mess up encryption status for toast table
+CREATE TABLE cities (
+  name       varchar(8),
+  population real,
+  elevation  int
+) USING tde_heap;
+CREATE TABLE state_capitals (
+  state      char(2) UNIQUE NOT NULL
+) INHERITS (cities) USING heap;
+CREATE TABLE capitals (
+  country      char(2) UNIQUE NOT NULL
+) INHERITS (cities) USING tde_heap;
+ALTER TABLE cities ALTER COLUMN name TYPE TEXT;
+ERROR:  Recursive ALTER TABLE on a mix of encrypted and unencrypted relations is not supported
+-- Enforce the same for typed tables
+CREATE TYPE people_type AS (age int, name varchar(8), dob date);
+CREATE TABLE sales_staff OF people_type USING tde_heap;
+CREATE TABLE other_staff OF people_type USING heap;
+ALTER TYPE people_type ALTER ATTRIBUTE name TYPE text CASCADE;
+ERROR:  Recursive ALTER TABLE on a mix of encrypted and unencrypted relations is not supported
+-- If all tpyed tables are encrypted everything should work as usual
+ALTER TABLE other_staff SET ACCESS METHOD tde_heap;
+ALTER TYPE people_type ALTER ATTRIBUTE name TYPE text CASCADE;
+SELECT pg_tde_is_encrypted('pg_toast.pg_toast_' || 'sales_staff'::regclass::oid);
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+SELECT pg_tde_is_encrypted('pg_toast.pg_toast_' || 'other_staff'::regclass::oid);
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+DROP TYPE people_type CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to table sales_staff
+drop cascades to table other_staff
+DROP TABLE cities CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to table state_capitals
+drop cascades to table capitals
 DROP TABLE plain_table;
 DROP EXTENSION pg_tde CASCADE;
 NOTICE:  drop cascades to 7 other objects

contrib/pg_tde/sql/partition_table.sql

@@ -31,6 +31,13 @@ SELECT pg_tde_is_encrypted('partition_q2_2024');
 SELECT pg_tde_is_encrypted('partition_q3_2024');
 SELECT pg_tde_is_encrypted('partition_q4_2024');
 
+-- Does not care about parent AM as long as all children with storage use the same
+ALTER TABLE partition_q1_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partition_q2_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partition_q3_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partition_q4_2024 SET ACCESS METHOD tde_heap;
+ALTER TABLE partitioned_table SET ACCESS METHOD heap;
+
 DROP TABLE partitioned_table;
 
 -- Partition inherits encryption status from parent table if default is heap and parent is tde_heap

contrib/pg_tde/sql/recreate_storage.sql

@@ -96,6 +96,39 @@ SELECT pg_tde_is_encrypted('plain_table_id2_seq2');
 ALTER SEQUENCE plain_table_id2_seq2 OWNED BY NONE;
 SELECT pg_tde_is_encrypted('plain_table_id2_seq2');
 
+-- Enforce that we do not mess up encryption status for toast table
+CREATE TABLE cities (
+  name       varchar(8),
+  population real,
+  elevation  int
+) USING tde_heap;
+
+CREATE TABLE state_capitals (
+  state      char(2) UNIQUE NOT NULL
+) INHERITS (cities) USING heap;
+
+CREATE TABLE capitals (
+  country      char(2) UNIQUE NOT NULL
+) INHERITS (cities) USING tde_heap;
+
+ALTER TABLE cities ALTER COLUMN name TYPE TEXT;
+
+-- Enforce the same for typed tables
+CREATE TYPE people_type AS (age int, name varchar(8), dob date);
+CREATE TABLE sales_staff OF people_type USING tde_heap;
+CREATE TABLE other_staff OF people_type USING heap;
+
+ALTER TYPE people_type ALTER ATTRIBUTE name TYPE text CASCADE;
+
+-- If all tpyed tables are encrypted everything should work as usual
+ALTER TABLE other_staff SET ACCESS METHOD tde_heap;
+ALTER TYPE people_type ALTER ATTRIBUTE name TYPE text CASCADE;
+
+SELECT pg_tde_is_encrypted('pg_toast.pg_toast_' || 'sales_staff'::regclass::oid);
+SELECT pg_tde_is_encrypted('pg_toast.pg_toast_' || 'other_staff'::regclass::oid);
+
+DROP TYPE people_type CASCADE;
+DROP TABLE cities CASCADE;
 DROP TABLE plain_table;
 DROP EXTENSION pg_tde CASCADE;
 RESET default_table_access_method;

contrib/pg_tde/src/pg_tde_event_capture.c

@@ -17,8 +17,10 @@
 #include "utils/lsyscache.h"
 #include "catalog/pg_class.h"
 #include "catalog/pg_database.h"
+#include "catalog/pg_inherits.h"
 #include "commands/defrem.h"
 #include "commands/sequence.h"
+#include "access/heapam.h"
 #include "access/table.h"
 #include "access/relation.h"
 #include "catalog/pg_event_trigger.h"
@@ -141,6 +143,97 @@ push_event_stack(const TdeDdlEvent *event)
 	MemoryContextSwitchTo(oldCtx);
 }
 
+static List *
+find_typed_table_dependencies(Oid typeOid)
+{
+	Relation	classRel;
+	ScanKeyData key[1];
+	TableScanDesc scan;
+	HeapTuple	tuple;
+	List	   *result = NIL;
+
+	classRel = table_open(RelationRelationId, AccessShareLock);
+
+	ScanKeyInit(&key[0],
+				Anum_pg_class_reloftype,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(typeOid));
+
+	scan = table_beginscan_catalog(classRel, 1, key);
+
+	while ((tuple = heap_getnext(scan, ForwardScanDirection)) != NULL)
+	{
+		Form_pg_class classform = (Form_pg_class) GETSTRUCT(tuple);
+
+		LockRelationOid(classform->oid, AccessShareLock);
+		result = lappend_oid(result, classform->oid);
+	}
+
+	table_endscan(scan);
+	table_close(classRel, AccessShareLock);
+
+	return result;
+}
+
+typedef enum
+{
+	ENC_MIX_UNKNOWN,
+	ENC_MIX_PLAIN,
+	ENC_MIX_ENCRYPTED,
+	ENC_MIX_MIXED,
+} EncryptionMix;
+
+/*
+ * Since ALTER TABLE can modify multiple tables due to inheritance or typed
+ * tables which can for example result in TOAST tables being created for some
+ * or all of the modified tables while the event trigger is only fired once we
+ * cannot rely on the event stack to make sure we get the correct encryption
+ * status.
+ *
+ * Our solution is to be cautious and only modify tables when all tables with
+ * storage are either encrypted or not encrypted. If there is a mix we will
+ * throw and error. The result of this is also used to properly inform the
+ * SMGR of the current encryption status.
+ */
+static EncryptionMix
+alter_table_encryption_mix(Oid relid)
+{
+	EncryptionMix enc = ENC_MIX_UNKNOWN;
+	Relation	rel;
+	List	   *children;
+	ListCell   *lc;
+
+	rel = relation_open(relid, NoLock);
+
+	if (RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
+		enc = rel->rd_rel->relam == get_tde_table_am_oid() ? ENC_MIX_ENCRYPTED : ENC_MIX_PLAIN;
+
+	if (rel->rd_rel->relkind == RELKIND_COMPOSITE_TYPE)
+		children = find_typed_table_dependencies(rel->rd_rel->reltype);
+	else
+		children = find_inheritance_children(relid, AccessShareLock);
+
+	relation_close(rel, NoLock);
+
+	foreach(lc, children)
+	{
+		Oid			childid = lfirst_oid(lc);
+		EncryptionMix childenc;
+
+		childenc = alter_table_encryption_mix(childid);
+
+		if (childenc != ENC_MIX_UNKNOWN)
+		{
+			if (enc == ENC_MIX_UNKNOWN)
+				enc = childenc;
+			else if (enc != childenc)
+				return ENC_MIX_MIXED;
+		}
+	}
+
+	return enc;
+}
+
 /*
  * pg_tde_ddl_command_start_capture is an event trigger function triggered
  * at the start of any DDL command execution.
@@ -267,6 +360,7 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 			AlterTableCmd *setAccessMethod = NULL;
 			ListCell   *lcmd;
 			TdeDdlEvent event = {.parsetree = parsetree};
+			EncryptionMix encmix;
 
 			foreach(lcmd, stmt->cmds)
 			{
@@ -276,6 +370,20 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 					setAccessMethod = cmd;
 			}
 
+			encmix = alter_table_encryption_mix(relid);
+
+			/*
+			 * This check is very braod and could be limited only to commands
+			 * which recurse to child tables or to those which may create new
+			 * relfilenodes, but this restrictive code is good enough for now.
+			 */
+			if (encmix == ENC_MIX_MIXED)
+			{
+				ereport(ERROR,
+						errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("Recursive ALTER TABLE on a mix of encrypted and unencrypted relations is not supported"));
+			}
+
 			/*
 			 * With a SET ACCESS METHOD clause, use that as the basis for
 			 * decisions. But if it's not present, look up encryption status
@@ -292,19 +400,13 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 			}
 			else
 			{
-				Relation	rel = relation_open(relid, AccessShareLock);
-
-				if (rel->rd_rel->relam == get_tde_table_am_oid())
+				if (encmix == ENC_MIX_ENCRYPTED)
 				{
-					/*
-					 * We are altering an encrypted table and ALTER TABLE can
-					 * possibly create new files so set the global state.
-					 */
 					event.encryptMode = TDE_ENCRYPT_MODE_ENCRYPT;
 					checkPrincipalKeyConfigured();
 				}
-
-				relation_close(rel, NoLock);
+				else if (encmix == ENC_MIX_PLAIN)
+					event.encryptMode = TDE_ENCRYPT_MODE_PLAIN;
 			}
 
 			push_event_stack(&event);