Fix the issue that PIVOT execution not return row affected line (babe…

…lfish-for-postgresql#2071) Previous our implement of PIVOT has a issue that it didn't return the row affected line as result, this fix has fixed the bug. Task: BABEL-284 Signed-off-by: Zhibai Song <[email protected]>
amazon-aurora · Dec 21, 2023 · 148635e · 148635e
1 parent 6b22380
commit 148635e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/contrib/babelfishpg_tsql/src/pl_exec.c b/contrib/babelfishpg_tsql/src/pl_exec.c
@@ -4693,6 +4693,12 @@ exec_stmt_execsql(PLtsql_execstate *estate,
 			return ret;
 		}
 
+		if (expr->plan && expr->plan->oneshot)
+		{
+			SPI_freeplan(expr->plan);
+			expr->plan = NULL;
+		}
+
 		if (expr->plan == NULL)
 		{
 			/*
@@ -5056,11 +5062,10 @@ exec_stmt_execsql(PLtsql_execstate *estate,
 		/* If query affects IDENTITY_INSERT relation then update sequence */
 		pltsql_update_identity_insert_sequence(expr);
 
-		/* If current plan constains a pivot operator, we remove the plan */
+		/* If current plan constains a pivot operator, we set it as execute oneshot */
 		if (is_pivot)
 		{
-			SPI_freeplan(expr->plan);
-			expr->plan = NULL;
+			expr->plan->oneshot = true;
 		}
 
 		/* Expect SPI_tuptable to be NULL else complain */

diff --git a/test/JDBC/expected/pivot-vu-verify.out b/test/JDBC/expected/pivot-vu-verify.out
@@ -738,6 +738,8 @@ PIVOT (
 ) AS pvt;
 SELECT TOP 10 * FROM pivot_insert_into ORDER by 1, 2;
 GO
+~~ROW COUNT: 197~~
+
 ~~START~~
 int#!#int#!#int#!#int#!#int#!#int#!#int
 1200#!#200#!#0#!#0#!#0#!#0#!#0