Assert TopTransactionContext existence instead of transaction state

If transaction was aborted for some reason there are still statements
that allowed to execute. So intead of assertion for transaction state
check if TopTransactionContext exists as it's supposed to be our parent
context.

Author: artemgavrilov

regression/expected/rollback.out

@@ -66,6 +66,53 @@ SELECT pg_stat_monitor_reset();
 
 (1 row)
 
+-- We still able to record ROLLBACK TO SAVEPOINT if subtransaction is aborted due an error.
+BEGIN;
+SAVEPOINT sp1;
+SELECT 1/0;
+ERROR:  division by zero
+ROLLBACK TO SAVEPOINT sp1;
+INSERT INTO t VALUES (3);
+COMMIT;
+SELECT query, username, datname FROM pg_stat_monitor ORDER BY query COLLATE "C";
+             query              | username |      datname       
+--------------------------------+----------+--------------------
+ BEGIN                          | u        | contrib_regression
+ COMMIT                         | u        | contrib_regression
+ INSERT INTO t VALUES (3)       | u        | contrib_regression
+ ROLLBACK TO SAVEPOINT sp1      |          | contrib_regression
+ SAVEPOINT sp1                  | u        | contrib_regression
+ SELECT 1/0;                    | u        | contrib_regression
+ SELECT pg_stat_monitor_reset() | u        | contrib_regression
+(7 rows)
+
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+-- We are still able to record PREPARE TRANSACTION issued in an aborted transaction, it rolls it back.
+BEGIN;
+INSERT INTO t VALUES (1);
+ERROR:  duplicate key value violates unique constraint "t_pkey"
+DETAIL:  Key (a)=(1) already exists.
+PREPARE TRANSACTION 'pgsm_aborted_transaction';
+SELECT query, username, datname FROM pg_stat_monitor ORDER BY query COLLATE "C";
+                     query                      | username |      datname       
+------------------------------------------------+----------+--------------------
+ BEGIN                                          | u        | contrib_regression
+ INSERT INTO t VALUES (1);                      | u        | contrib_regression
+ PREPARE TRANSACTION 'pgsm_aborted_transaction' |          | contrib_regression
+ SELECT pg_stat_monitor_reset()                 | u        | contrib_regression
+(4 rows)
+
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
 DROP TABLE t;
 SET ROLE NONE;
 DROP USER u;

regression/sql/rollback.sql

@@ -29,6 +29,25 @@ ROLLBACK;
 SELECT query, username, datname FROM pg_stat_monitor ORDER BY query COLLATE "C";
 SELECT pg_stat_monitor_reset();
 
+-- We still able to record ROLLBACK TO SAVEPOINT if subtransaction is aborted due an error.
+BEGIN;
+SAVEPOINT sp1;
+SELECT 1/0;
+ROLLBACK TO SAVEPOINT sp1;
+INSERT INTO t VALUES (3);
+COMMIT;
+
+SELECT query, username, datname FROM pg_stat_monitor ORDER BY query COLLATE "C";
+SELECT pg_stat_monitor_reset();
+
+-- We are still able to record PREPARE TRANSACTION issued in an aborted transaction, it rolls it back.
+BEGIN;
+INSERT INTO t VALUES (1);
+PREPARE TRANSACTION 'pgsm_aborted_transaction';
+
+SELECT query, username, datname FROM pg_stat_monitor ORDER BY query COLLATE "C";
+SELECT pg_stat_monitor_reset();
+
 DROP TABLE t;
 SET ROLE NONE;
 DROP USER u;

src/pg_stat_monitor.c

@@ -1588,7 +1588,7 @@ pgsm_store_error(const char *query, const ErrorData *edata)
 static MemoryContext
 pgsm_memory_context(void)
 {
-	Assert(IsTransactionState());
+	Assert(TopTransactionContext != NULL);
 
 	if (PgsmMemoryContext == NULL)
 	{
@@ -1597,8 +1597,8 @@ pgsm_memory_context(void)
 		 * scenario. CurrentMemoryContext here is just a failsafe mechanism,
 		 * it should never happen.
 		 */
-		MemoryContext parent = IsTransactionState() ? TopTransactionContext
-			: CurrentMemoryContext;
+		MemoryContext parent = TopTransactionContext != NULL
+			? TopTransactionContext : CurrentMemoryContext;
 
 		PgsmMemoryContext = AllocSetContextCreate(parent,
 												  "pg_stat_monitor local store",