From dffece72a53f0fd3b470642babc97d6558b0e195 Mon Sep 17 00:00:00 2001
From: Andreas Karlsson <andreas.karlsson@percona.com>
Date: Thu, 24 Apr 2025 14:44:20 +0200
Subject: [PATCH] Make ALTER TABLE ... SET ACCESS METHOD logic easier to read

Also add a couple of tests for the DEFAULT case to avoid regressions.
---
 .../pg_tde/expected/change_access_method.out  | 18 +++++++++
 contrib/pg_tde/sql/change_access_method.sql   | 15 ++++++++
 contrib/pg_tde/src/pg_tde_event_capture.c     | 37 +++++++++----------
 3 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/contrib/pg_tde/expected/change_access_method.out b/contrib/pg_tde/expected/change_access_method.out
index 886bc857373..ce8f54af706 100644
--- a/contrib/pg_tde/expected/change_access_method.out
+++ b/contrib/pg_tde/expected/change_access_method.out
@@ -121,6 +121,24 @@ SELECT pg_tde_is_encrypted('country_table_pkey');
  t
 (1 row)
 
+-- Test that we honor the default value
+SET default_table_access_method = 'heap';
+ALTER TABLE country_table SET ACCESS METHOD DEFAULT;
+SELECT pg_tde_is_encrypted('country_table');
+ pg_tde_is_encrypted 
+---------------------
+ f
+(1 row)
+
+SET default_table_access_method = 'tde_heap';
+ALTER TABLE country_table SET ACCESS METHOD DEFAULT;
+SELECT pg_tde_is_encrypted('country_table');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+RESET default_table_access_method;
 ALTER TABLE country_table ADD y text;
 SELECT pg_tde_is_encrypted('pg_toast.pg_toast_' || 'country_table'::regclass::oid);
  pg_tde_is_encrypted 
diff --git a/contrib/pg_tde/sql/change_access_method.sql b/contrib/pg_tde/sql/change_access_method.sql
index 3e656a5afd4..dc01fbed895 100644
--- a/contrib/pg_tde/sql/change_access_method.sql
+++ b/contrib/pg_tde/sql/change_access_method.sql
@@ -46,6 +46,21 @@ SELECT pg_tde_is_encrypted('country_table');
 SELECT pg_tde_is_encrypted('country_table_country_id_seq');
 SELECT pg_tde_is_encrypted('country_table_pkey');
 
+-- Test that we honor the default value
+SET default_table_access_method = 'heap';
+
+ALTER TABLE country_table SET ACCESS METHOD DEFAULT;
+
+SELECT pg_tde_is_encrypted('country_table');
+
+SET default_table_access_method = 'tde_heap';
+
+ALTER TABLE country_table SET ACCESS METHOD DEFAULT;
+
+SELECT pg_tde_is_encrypted('country_table');
+
+RESET default_table_access_method;
+
 ALTER TABLE country_table ADD y text;
 
 SELECT pg_tde_is_encrypted('pg_toast.pg_toast_' || 'country_table'::regclass::oid);
diff --git a/contrib/pg_tde/src/pg_tde_event_capture.c b/contrib/pg_tde/src/pg_tde_event_capture.c
index d9cbe1ebf75..d6b7a7ab4bb 100644
--- a/contrib/pg_tde/src/pg_tde_event_capture.c
+++ b/contrib/pg_tde/src/pg_tde_event_capture.c
@@ -32,7 +32,6 @@
 
 /* Global variable that gets set at ddl start and cleard out at ddl end*/
 static TdeCreateEvent tdeCurrentCreateEvent = {.tid = {.value = 0}};
-static bool alterSetAccessMethod = false;
 
 static void reset_current_tde_create_event(void);
 static Oid	get_tde_table_am_oid(void);
@@ -176,6 +175,7 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 		AlterTableStmt *stmt = castNode(AlterTableStmt, parsetree);
 		ListCell   *lcmd;
 		Oid			relationId = RangeVarGetRelid(stmt->relation, AccessShareLock, true);
+		AlterTableCmd *setAccessMethod = NULL;
 
 		validateCurrentEventTriggerState(true);
 		tdeCurrentCreateEvent.tid = GetCurrentFullTransactionId();
@@ -185,29 +185,27 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 			AlterTableCmd *cmd = castNode(AlterTableCmd, lfirst(lcmd));
 
 			if (cmd->subtype == AT_SetAccessMethod)
-			{
-				tdeCurrentCreateEvent.baseTableOid = relationId;
-				tdeCurrentCreateEvent.alterAccessMethodMode = true;
-
-				if (shouldEncryptTable(cmd->name))
-					tdeCurrentCreateEvent.encryptMode = true;
-
-				checkEncryptionStatus();
-
-				alterSetAccessMethod = true;
-			}
+				setAccessMethod = cmd;
 		}
 
-		if (!alterSetAccessMethod)
+		tdeCurrentCreateEvent.baseTableOid = relationId;
+
+		/*
+		 * With a SET ACCESS METHOD clause, use that as the basis for
+		 * decisions. But if it's not present, look up encryption status of
+		 * the table.
+		 */
+		if (setAccessMethod)
 		{
-			/*
-			 * With a SET ACCESS METHOD clause, use that as the basis for
-			 * decisions. But if it's not present, look up encryption status
-			 * of the table.
-			 */
+			if (shouldEncryptTable(setAccessMethod->name))
+				tdeCurrentCreateEvent.encryptMode = true;
 
-			tdeCurrentCreateEvent.baseTableOid = relationId;
+			checkEncryptionStatus();
 
+			tdeCurrentCreateEvent.alterAccessMethodMode = true;
+		}
+		else
+		{
 			if (relationId != InvalidOid)
 			{
 				Relation	rel = relation_open(relationId, NoLock);
@@ -307,7 +305,6 @@ reset_current_tde_create_event(void)
 	tdeCurrentCreateEvent.baseTableOid = InvalidOid;
 	tdeCurrentCreateEvent.tid = InvalidFullTransactionId;
 	tdeCurrentCreateEvent.alterAccessMethodMode = false;
-	alterSetAccessMethod = false;
 }
 
 static Oid