PL/pgSQL: Add support for SET TRANSACTION

A normal SQL command run inside PL/pgSQL acquires a snapshot, but SET
TRANSACTION does not work anymore if a snapshot is set.  So we have to
handle this separately.

Reviewed-by: Alexander Korotkov <a.korotkov@postgrespro.ru>
Reviewed-by: Tomas Vondra <tomas.vondra@2ndquadrant.com>

7 files changed, 157 insertions, 2 deletions

diff --git a/src/pl/plpgsql/src/expected/plpgsql_transaction.out b/src/pl/plpgsql/src/expected/plpgsql_transaction.out
index b7f77101c3d..2d0e3fa85ed 100644
--- a/src/pl/plpgsql/src/expected/plpgsql_transaction.out
+++ b/src/pl/plpgsql/src/expected/plpgsql_transaction.out
@@ -389,6 +389,35 @@ SELECT * FROM test3;
  1
 (1 row)
 
+-- SET TRANSACTION
+DO LANGUAGE plpgsql $$
+BEGIN
+    PERFORM 1;
+    RAISE INFO '%', current_setting('transaction_isolation');
+    COMMIT;
+    SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+    PERFORM 1;
+    RAISE INFO '%', current_setting('transaction_isolation');
+    COMMIT;
+    SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+    RESET TRANSACTION ISOLATION LEVEL;
+    PERFORM 1;
+    RAISE INFO '%', current_setting('transaction_isolation');
+    COMMIT;
+END;
+$$;
+INFO:  read committed
+INFO:  repeatable read
+INFO:  read committed
+-- error case
+DO LANGUAGE plpgsql $$
+BEGIN
+    SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+END;
+$$;
+ERROR:  SET TRANSACTION ISOLATION LEVEL must be called before any query
+CONTEXT:  SQL statement "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
+PL/pgSQL function inline_code_block line 3 at SET
 DROP TABLE test1;
 DROP TABLE test2;
 DROP TABLE test3;
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 255bdbf2c8e..4b6f51ac98d 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -305,6 +305,8 @@ static int exec_stmt_commit(PLpgSQL_execstate *estate,
 				 PLpgSQL_stmt_commit *stmt);
 static int exec_stmt_rollback(PLpgSQL_execstate *estate,
 				   PLpgSQL_stmt_rollback *stmt);
+static int exec_stmt_set(PLpgSQL_execstate *estate,
+				   PLpgSQL_stmt_set *stmt);
 
 static void plpgsql_estate_setup(PLpgSQL_execstate *estate,
 					 PLpgSQL_function *func,
@@ -2005,6 +2007,10 @@ exec_stmt(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 			rc = exec_stmt_rollback(estate, (PLpgSQL_stmt_rollback *) stmt);
 			break;
 
+		case PLPGSQL_STMT_SET:
+			rc = exec_stmt_set(estate, (PLpgSQL_stmt_set *) stmt);
+			break;
+
 		default:
 			estate->err_stmt = save_estmt;
 			elog(ERROR, "unrecognized cmd_type: %d", stmt->cmd_type);
@@ -4732,6 +4738,35 @@ exec_stmt_rollback(PLpgSQL_execstate *estate, PLpgSQL_stmt_rollback *stmt)
 	return PLPGSQL_RC_OK;
 }
 
+/*
+ * exec_stmt_set
+ *
+ * Execute SET/RESET statement.
+ *
+ * We just parse and execute the statement normally, but we have to do it
+ * without setting a snapshot, for things like SET TRANSACTION.
+ */
+static int
+exec_stmt_set(PLpgSQL_execstate *estate, PLpgSQL_stmt_set *stmt)
+{
+	PLpgSQL_expr *expr = stmt->expr;
+	int			rc;
+
+	if (expr->plan == NULL)
+	{
+		exec_prepare_plan(estate, expr, 0, true);
+		expr->plan->no_snapshots = true;
+	}
+
+	rc = SPI_execute_plan(expr->plan, NULL, NULL, estate->readonly_func, 0);
+
+	if (rc != SPI_OK_UTILITY)
+		elog(ERROR, "SPI_execute_plan failed executing query \"%s\": %s",
+			 expr->query, SPI_result_code_string(rc));
+
+	return PLPGSQL_RC_OK;
+}
+
 /* ----------
  * exec_assign_expr			Put an expression's result into a variable.
  * ----------
diff --git a/src/pl/plpgsql/src/pl_funcs.c b/src/pl/plpgsql/src/pl_funcs.c
index fc96fb5f4d1..b93f8662233 100644
--- a/src/pl/plpgsql/src/pl_funcs.c
+++ b/src/pl/plpgsql/src/pl_funcs.c
@@ -290,6 +290,8 @@ plpgsql_stmt_typename(PLpgSQL_stmt *stmt)
 			return "COMMIT";
 		case PLPGSQL_STMT_ROLLBACK:
 			return "ROLLBACK";
+		case PLPGSQL_STMT_SET:
+			return "SET";
 	}
 
 	return "unknown";
@@ -372,6 +374,7 @@ static void free_perform(PLpgSQL_stmt_perform *stmt);
 static void free_call(PLpgSQL_stmt_call *stmt);
 static void free_commit(PLpgSQL_stmt_commit *stmt);
 static void free_rollback(PLpgSQL_stmt_rollback *stmt);
+static void free_set(PLpgSQL_stmt_set *stmt);
 static void free_expr(PLpgSQL_expr *expr);
 
 
@@ -461,6 +464,9 @@ free_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_ROLLBACK:
 			free_rollback((PLpgSQL_stmt_rollback *) stmt);
 			break;
+		case PLPGSQL_STMT_SET:
+			free_set((PLpgSQL_stmt_set *) stmt);
+			break;
 		default:
 			elog(ERROR, "unrecognized cmd_type: %d", stmt->cmd_type);
 			break;
@@ -625,6 +631,12 @@ free_rollback(PLpgSQL_stmt_rollback *stmt)
 }
 
 static void
+free_set(PLpgSQL_stmt_set *stmt)
+{
+	free_expr(stmt->expr);
+}
+
+static void
 free_exit(PLpgSQL_stmt_exit *stmt)
 {
 	free_expr(stmt->cond);
@@ -820,6 +832,7 @@ static void dump_perform(PLpgSQL_stmt_perform *stmt);
 static void dump_call(PLpgSQL_stmt_call *stmt);
 static void dump_commit(PLpgSQL_stmt_commit *stmt);
 static void dump_rollback(PLpgSQL_stmt_rollback *stmt);
+static void dump_set(PLpgSQL_stmt_set *stmt);
 static void dump_expr(PLpgSQL_expr *expr);
 
 
@@ -919,6 +932,9 @@ dump_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_ROLLBACK:
 			dump_rollback((PLpgSQL_stmt_rollback *) stmt);
 			break;
+		case PLPGSQL_STMT_SET:
+			dump_set((PLpgSQL_stmt_set *) stmt);
+			break;
 		default:
 			elog(ERROR, "unrecognized cmd_type: %d", stmt->cmd_type);
 			break;
@@ -1315,6 +1331,13 @@ dump_rollback(PLpgSQL_stmt_rollback *stmt)
 }
 
 static void
+dump_set(PLpgSQL_stmt_set *stmt)
+{
+	dump_ind();
+	printf("%s\n", stmt->expr->query);
+}
+
+static void
 dump_exit(PLpgSQL_stmt_exit *stmt)
 {
 	dump_ind();
diff --git a/src/pl/plpgsql/src/pl_gram.y b/src/pl/plpgsql/src/pl_gram.y
index d928efa13ea..f9ba19cbdf0 100644
--- a/src/pl/plpgsql/src/pl_gram.y
+++ b/src/pl/plpgsql/src/pl_gram.y
@@ -199,7 +199,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %type <stmt>	stmt_return stmt_raise stmt_assert stmt_execsql
 %type <stmt>	stmt_dynexecute stmt_for stmt_perform stmt_call stmt_getdiag
 %type <stmt>	stmt_open stmt_fetch stmt_move stmt_close stmt_null
-%type <stmt>	stmt_commit stmt_rollback
+%type <stmt>	stmt_commit stmt_rollback stmt_set
 %type <stmt>	stmt_case stmt_foreach_a
 
 %type <list>	proc_exceptions
@@ -327,6 +327,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %token <keyword>	K_QUERY
 %token <keyword>	K_RAISE
 %token <keyword>	K_RELATIVE
+%token <keyword>	K_RESET
 %token <keyword>	K_RESULT_OID
 %token <keyword>	K_RETURN
 %token <keyword>	K_RETURNED_SQLSTATE
@@ -337,6 +338,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %token <keyword>	K_SCHEMA
 %token <keyword>	K_SCHEMA_NAME
 %token <keyword>	K_SCROLL
+%token <keyword>	K_SET
 %token <keyword>	K_SLICE
 %token <keyword>	K_SQLSTATE
 %token <keyword>	K_STACKED
@@ -893,6 +895,8 @@ proc_stmt		: pl_block ';'
 						{ $$ = $1; }
 				| stmt_rollback
 						{ $$ = $1; }
+				| stmt_set
+						{ $$ = $1; }
 				;
 
 stmt_perform	: K_PERFORM expr_until_semi
@@ -2206,6 +2210,30 @@ stmt_rollback	: K_ROLLBACK ';'
 					}
 				;
 
+stmt_set	: K_SET
+					{
+						PLpgSQL_stmt_set *new;
+
+						new = palloc0(sizeof(PLpgSQL_stmt_set));
+						new->cmd_type = PLPGSQL_STMT_SET;
+						new->lineno = plpgsql_location_to_lineno(@1);
+						new->expr = read_sql_stmt("SET ");
+
+						$$ = (PLpgSQL_stmt *)new;
+					}
+			| K_RESET
+					{
+						PLpgSQL_stmt_set *new;
+
+						new = palloc0(sizeof(PLpgSQL_stmt_set));
+						new->cmd_type = PLPGSQL_STMT_SET;
+						new->lineno = plpgsql_location_to_lineno(@1);
+						new->expr = read_sql_stmt("RESET ");
+
+						$$ = (PLpgSQL_stmt *)new;
+					}
+			;
+
 
 cursor_variable	: T_DATUM
 					{
@@ -2494,6 +2522,7 @@ unreserved_keyword	:
 				| K_QUERY
 				| K_RAISE
 				| K_RELATIVE
+				| K_RESET
 				| K_RESULT_OID
 				| K_RETURN
 				| K_RETURNED_SQLSTATE
@@ -2504,6 +2533,7 @@ unreserved_keyword	:
 				| K_SCHEMA
 				| K_SCHEMA_NAME
 				| K_SCROLL
+				| K_SET
 				| K_SLICE
 				| K_SQLSTATE
 				| K_STACKED
diff --git a/src/pl/plpgsql/src/pl_scanner.c b/src/pl/plpgsql/src/pl_scanner.c
index 81b5bd1fe03..256fc0a243b 100644
--- a/src/pl/plpgsql/src/pl_scanner.c
+++ b/src/pl/plpgsql/src/pl_scanner.c
@@ -158,6 +158,7 @@ static const ScanKeyword unreserved_keywords[] = {
 	PG_KEYWORD("query", K_QUERY, UNRESERVED_KEYWORD)
 	PG_KEYWORD("raise", K_RAISE, UNRESERVED_KEYWORD)
 	PG_KEYWORD("relative", K_RELATIVE, UNRESERVED_KEYWORD)
+	PG_KEYWORD("reset", K_RESET, UNRESERVED_KEYWORD)
 	PG_KEYWORD("result_oid", K_RESULT_OID, UNRESERVED_KEYWORD)
 	PG_KEYWORD("return", K_RETURN, UNRESERVED_KEYWORD)
 	PG_KEYWORD("returned_sqlstate", K_RETURNED_SQLSTATE, UNRESERVED_KEYWORD)
@@ -168,6 +169,7 @@ static const ScanKeyword unreserved_keywords[] = {
 	PG_KEYWORD("schema", K_SCHEMA, UNRESERVED_KEYWORD)
 	PG_KEYWORD("schema_name", K_SCHEMA_NAME, UNRESERVED_KEYWORD)
 	PG_KEYWORD("scroll", K_SCROLL, UNRESERVED_KEYWORD)
+	PG_KEYWORD("set", K_SET, UNRESERVED_KEYWORD)
 	PG_KEYWORD("slice", K_SLICE, UNRESERVED_KEYWORD)
 	PG_KEYWORD("sqlstate", K_SQLSTATE, UNRESERVED_KEYWORD)
 	PG_KEYWORD("stacked", K_STACKED, UNRESERVED_KEYWORD)
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 3f617201d67..8d30180d424 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -127,7 +127,8 @@ typedef enum PLpgSQL_stmt_type
 	PLPGSQL_STMT_PERFORM,
 	PLPGSQL_STMT_CALL,
 	PLPGSQL_STMT_COMMIT,
-	PLPGSQL_STMT_ROLLBACK
+	PLPGSQL_STMT_ROLLBACK,
+	PLPGSQL_STMT_SET
 } PLpgSQL_stmt_type;
 
 /*
@@ -540,6 +541,16 @@ typedef struct PLpgSQL_stmt_rollback
 } PLpgSQL_stmt_rollback;
 
 /*
+ * SET statement
+ */
+typedef struct PLpgSQL_stmt_set
+{
+	PLpgSQL_stmt_type cmd_type;
+	int			lineno;
+	PLpgSQL_expr *expr;
+} PLpgSQL_stmt_set;
+
+/*
  * GET DIAGNOSTICS item
  */
 typedef struct PLpgSQL_diag_item
diff --git a/src/pl/plpgsql/src/sql/plpgsql_transaction.sql b/src/pl/plpgsql/src/sql/plpgsql_transaction.sql
index b7b6f81128a..373d89864a2 100644
--- a/src/pl/plpgsql/src/sql/plpgsql_transaction.sql
+++ b/src/pl/plpgsql/src/sql/plpgsql_transaction.sql
@@ -317,6 +317,31 @@ $$;
 SELECT * FROM test3;
 
 
+-- SET TRANSACTION
+DO LANGUAGE plpgsql $$
+BEGIN
+    PERFORM 1;
+    RAISE INFO '%', current_setting('transaction_isolation');
+    COMMIT;
+    SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+    PERFORM 1;
+    RAISE INFO '%', current_setting('transaction_isolation');
+    COMMIT;
+    SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+    RESET TRANSACTION ISOLATION LEVEL;
+    PERFORM 1;
+    RAISE INFO '%', current_setting('transaction_isolation');
+    COMMIT;
+END;
+$$;
+
+-- error case
+DO LANGUAGE plpgsql $$
+BEGIN
+    SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+END;
+$$;
+
 DROP TABLE test1;
 DROP TABLE test2;
 DROP TABLE test3;