v1-0001-trying-to-save-mt_transition_capture-while-ExecIn.patch

From cda2f8324ef920c807008e2763f97c6503c17a94 Mon Sep 17 00:00:00 2001
From: pgaddict <jian.universality@gmail.com>
Date: Wed, 6 Dec 2023 20:58:08 +0800
Subject: [PATCH v1 1/1] trying to save mt_transition_capture while ExecInsert

---
 src/backend/executor/nodeModifyTable.c       | 21 ++++-
 src/test/regress/expected/for_portion_of.out | 80 ++++++++++++++++++++
 src/test/regress/sql/for_portion_of.sql      | 72 ++++++++++++++++++
 3 files changed, 171 insertions(+), 2 deletions(-)

diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index fe57e592..c9d14dab 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -167,7 +167,8 @@ static bool ExecMergeMatched(ModifyTableContext *context,
 static void ExecMergeNotMatched(ModifyTableContext *context,
 								ResultRelInfo *resultRelInfo,
 								bool canSetTag);
-
+static void
+ExecSetupTransitionCaptureState(ModifyTableState *mtstate, EState *estate);

 /*
  * Verify that the tuples to be produced by INSERT match the
@@ -1259,7 +1260,7 @@ ExecForPortionOfLeftovers(ModifyTableContext *context,
 	TupleTableSlot *leftoverTuple2 = fpoState->fp_Leftover2;
 	HeapTuple oldtuple = NULL;
 	bool shouldFree = false;
-
+	int		i;	/* original cmd type */
 	/*
 	 * Get the range of the old pre-UPDATE/DELETE tuple,
 	 * so we can intersect it with the FOR PORTION OF target
@@ -1312,6 +1313,7 @@ ExecForPortionOfLeftovers(ModifyTableContext *context,
 	 * TODO: Really? What if you update the partition key?
 	 */

+	i	= context->mtstate->operation;
 	if (!RangeIsEmpty(leftoverRangeType1))
 	{
 		oldtuple = ExecFetchSlotHeapTuple(oldtupleSlot, false, &shouldFree);
@@ -1320,6 +1322,10 @@ ExecForPortionOfLeftovers(ModifyTableContext *context,
 		set_leftover_tuple_bounds(leftoverTuple1, forPortionOf, typcache, leftoverRangeType1);
 		ExecMaterializeSlot(leftoverTuple1);

+		context->mtstate->operation = CMD_INSERT;
+		context->mtstate->mt_transition_capture = NULL;
+		ExecSetupTransitionCaptureState(context->mtstate, estate);
+
 		// TODO: Need to save context->mtstate->mt_transition_capture? (See comment on ExecInsert)
 		ExecInsert(context, resultRelInfo, leftoverTuple1, node->canSetTag, NULL, NULL);
 	}
@@ -1340,10 +1346,21 @@ ExecForPortionOfLeftovers(ModifyTableContext *context,
 		set_leftover_tuple_bounds(leftoverTuple2, forPortionOf, typcache, leftoverRangeType2);
 		ExecMaterializeSlot(leftoverTuple2);

+		context->mtstate->operation = CMD_INSERT;
+		context->mtstate->mt_transition_capture = NULL;
+		ExecSetupTransitionCaptureState(context->mtstate, estate);
+
 		// TODO: Need to save context->mtstate->mt_transition_capture? (See comment on ExecInsert)
 		ExecInsert(context, resultRelInfo, leftoverTuple2, node->canSetTag, NULL, NULL);
 	}

+	if (!RangeIsEmpty(leftoverRangeType2) || !RangeIsEmpty(leftoverRangeType1))
+	{
+		/* if any of the above branch executed then */
+		context->mtstate->operation = i;
+		context->mtstate->mt_transition_capture = NULL;
+		ExecSetupTransitionCaptureState(context->mtstate, estate);
+	}
 	if (shouldFree)
 		heap_freetuple(oldtuple);
 }
diff --git a/src/test/regress/expected/for_portion_of.out b/src/test/regress/expected/for_portion_of.out
index 64f7bc7e..04202aa0 100644
--- a/src/test/regress/expected/for_portion_of.out
+++ b/src/test/regress/expected/for_portion_of.out
@@ -441,3 +441,83 @@ SELECT * FROM for_portion_of_test2 ORDER BY id, valid_at;

 DROP TABLE for_portion_of_test2;
 DROP TYPE mydaterange;
+DROP TABLE IF EXISTS for_portion_of_test;
+CREATE TABLE for_portion_of_test (id int4range,valid_at daterange, name text);
+INSERT INTO for_portion_of_test VALUES ('[1,1]', '[2018-01-01,2020-01-01)', 'one');
+create or replace function dump_trigger() returns trigger language plpgsql as
+$$
+	BEGIN
+    raise notice 'TG_OP = %', TG_OP;
+	IF TG_OP = 'INSERT' then
+		raise notice 'trigger = %, TG_LEVEL %, new table = %',
+					TG_NAME, TG_LEVEL,
+					(select string_agg(new_table::text, ', ' order by id) from new_table);
+    ELSIF TG_OP = 'UPDATE' then
+		raise notice 'trigger = %, TG_LEVEL %, old table = %, new table = %',
+					TG_NAME, TG_LEVEL,
+					(select string_agg(old_table::text, ', ' order by id) from old_table),
+					(select string_agg(new_table::text, ', ' order by id) from new_table);
+    ELSIF TG_OP = 'DELETE' then
+			raise notice 'trigger = %,TG_LEVEL %, old table = %',
+					TG_NAME, TG_LEVEL,
+					(select string_agg(old_table::text, ', ' order by id) from old_table);
+	END IF;
+    RETURN NULL;
+	END;
+$$;
+create or replace trigger for_portion_of_test_insert_trig
+after insert on for_portion_of_test
+referencing new table as new_table
+for each row execute procedure dump_trigger();
+create or replace trigger for_portion_of_test_insert_trig_stmt
+after insert on for_portion_of_test
+referencing new table as new_table
+for each statement execute procedure dump_trigger();
+create or replace trigger for_portion_of_test_update_trig
+after update on for_portion_of_test
+referencing old table as old_table new table as new_table
+for each row execute procedure dump_trigger();
+create or replace trigger for_portion_of_test_update_trig_stmt
+after update on for_portion_of_test
+referencing old table as old_table new table as new_table
+for each statement execute procedure dump_trigger();
+create or replace trigger for_portion_of_test_del_trig
+after delete on for_portion_of_test
+referencing old table as old_table
+for each row execute procedure dump_trigger();
+create or replace trigger for_portion_of_test_del_trig_stmt
+after delete on for_portion_of_test
+referencing old table as old_table
+for each statement execute procedure dump_trigger();
+begin;
+UPDATE for_portion_of_test FOR PORTION OF valid_at
+	FROM '2018-01-15' TO '2019-01-01' SET name = '2018-01-15_to_2019-01-01';
+NOTICE:  TG_OP = INSERT
+NOTICE:  trigger = for_portion_of_test_insert_trig, TG_LEVEL ROW, new table = ("[1,2)","[01-01-2018,01-15-2018)",one), ("[1,2)","[01-01-2019,01-01-2020)",one)
+NOTICE:  TG_OP = INSERT
+NOTICE:  trigger = for_portion_of_test_insert_trig, TG_LEVEL ROW, new table = ("[1,2)","[01-01-2018,01-15-2018)",one), ("[1,2)","[01-01-2019,01-01-2020)",one)
+NOTICE:  TG_OP = UPDATE
+NOTICE:  trigger = for_portion_of_test_update_trig, TG_LEVEL ROW, old table = ("[1,2)","[01-01-2018,01-01-2020)",one), new table = ("[1,2)","[01-15-2018,01-01-2019)",2018-01-15_to_2019-01-01)
+NOTICE:  TG_OP = UPDATE
+NOTICE:  trigger = for_portion_of_test_update_trig_stmt, TG_LEVEL STATEMENT, old table = ("[1,2)","[01-01-2018,01-01-2020)",one), new table = ("[1,2)","[01-15-2018,01-01-2019)",2018-01-15_to_2019-01-01)
+rollback;
+begin;
+delete from for_portion_of_test FOR PORTION OF valid_at FROM NULL TO '2018-01-21';
+NOTICE:  TG_OP = INSERT
+NOTICE:  trigger = for_portion_of_test_insert_trig, TG_LEVEL ROW, new table = ("[1,2)","[01-21-2018,01-01-2020)",one)
+NOTICE:  TG_OP = DELETE
+NOTICE:  trigger = for_portion_of_test_del_trig,TG_LEVEL ROW, old table = ("[1,2)","[01-01-2018,01-01-2020)",one)
+NOTICE:  TG_OP = DELETE
+NOTICE:  trigger = for_portion_of_test_del_trig_stmt,TG_LEVEL STATEMENT, old table = ("[1,2)","[01-01-2018,01-01-2020)",one)
+rollback;
+begin;
+UPDATE for_portion_of_test FOR PORTION OF valid_at
+	FROM NULL TO '2018-01-02' SET name = 'NULL_to_2018-01-01';
+NOTICE:  TG_OP = INSERT
+NOTICE:  trigger = for_portion_of_test_insert_trig, TG_LEVEL ROW, new table = ("[1,2)","[01-02-2018,01-01-2020)",one)
+NOTICE:  TG_OP = UPDATE
+NOTICE:  trigger = for_portion_of_test_update_trig, TG_LEVEL ROW, old table = ("[1,2)","[01-01-2018,01-01-2020)",one), new table = ("[1,2)","[01-01-2018,01-02-2018)",NULL_to_2018-01-01)
+NOTICE:  TG_OP = UPDATE
+NOTICE:  trigger = for_portion_of_test_update_trig_stmt, TG_LEVEL STATEMENT, old table = ("[1,2)","[01-01-2018,01-01-2020)",one), new table = ("[1,2)","[01-01-2018,01-02-2018)",NULL_to_2018-01-01)
+rollback;
+DROP TABLE IF EXISTS for_portion_of_test;
diff --git a/src/test/regress/sql/for_portion_of.sql b/src/test/regress/sql/for_portion_of.sql
index 677b8e38..e86819c4 100644
--- a/src/test/regress/sql/for_portion_of.sql
+++ b/src/test/regress/sql/for_portion_of.sql
@@ -341,3 +341,75 @@ SELECT * FROM for_portion_of_test2 ORDER BY id, valid_at;

 DROP TABLE for_portion_of_test2;
 DROP TYPE mydaterange;
+
+DROP TABLE IF EXISTS for_portion_of_test;
+CREATE TABLE for_portion_of_test (id int4range,valid_at daterange, name text);
+INSERT INTO for_portion_of_test VALUES ('[1,1]', '[2018-01-01,2020-01-01)', 'one');
+
+create or replace function dump_trigger() returns trigger language plpgsql as
+$$
+	BEGIN
+    raise notice 'TG_OP = %', TG_OP;
+	IF TG_OP = 'INSERT' then
+		raise notice 'trigger = %, TG_LEVEL %, new table = %',
+					TG_NAME, TG_LEVEL,
+					(select string_agg(new_table::text, ', ' order by id) from new_table);
+    ELSIF TG_OP = 'UPDATE' then
+		raise notice 'trigger = %, TG_LEVEL %, old table = %, new table = %',
+					TG_NAME, TG_LEVEL,
+					(select string_agg(old_table::text, ', ' order by id) from old_table),
+					(select string_agg(new_table::text, ', ' order by id) from new_table);
+    ELSIF TG_OP = 'DELETE' then
+			raise notice 'trigger = %,TG_LEVEL %, old table = %',
+					TG_NAME, TG_LEVEL,
+					(select string_agg(old_table::text, ', ' order by id) from old_table);
+	END IF;
+    RETURN NULL;
+	END;
+$$;
+
+create or replace trigger for_portion_of_test_insert_trig
+after insert on for_portion_of_test
+referencing new table as new_table
+for each row execute procedure dump_trigger();
+
+create or replace trigger for_portion_of_test_insert_trig_stmt
+after insert on for_portion_of_test
+referencing new table as new_table
+for each statement execute procedure dump_trigger();
+
+create or replace trigger for_portion_of_test_update_trig
+after update on for_portion_of_test
+referencing old table as old_table new table as new_table
+for each row execute procedure dump_trigger();
+
+create or replace trigger for_portion_of_test_update_trig_stmt
+after update on for_portion_of_test
+referencing old table as old_table new table as new_table
+for each statement execute procedure dump_trigger();
+
+create or replace trigger for_portion_of_test_del_trig
+after delete on for_portion_of_test
+referencing old table as old_table
+for each row execute procedure dump_trigger();
+
+create or replace trigger for_portion_of_test_del_trig_stmt
+after delete on for_portion_of_test
+referencing old table as old_table
+for each statement execute procedure dump_trigger();
+
+begin;
+UPDATE for_portion_of_test FOR PORTION OF valid_at
+	FROM '2018-01-15' TO '2019-01-01' SET name = '2018-01-15_to_2019-01-01';
+rollback;
+
+begin;
+delete from for_portion_of_test FOR PORTION OF valid_at FROM NULL TO '2018-01-21';
+rollback;
+
+begin;
+UPDATE for_portion_of_test FOR PORTION OF valid_at
+	FROM NULL TO '2018-01-02' SET name = 'NULL_to_2018-01-01';
+rollback;
+
+DROP TABLE IF EXISTS for_portion_of_test;
\ No newline at end of file
--
2.34.1