Add support for sys.sp_reset_connection stored procedure (babelfish-f…

…or-postgresql#2758)

Support sys.sp_reset_connection store procedure to reset tds connection

1. At the moment babelfish doesn't support sys.sp_reset_connection to reset the connection state. With this change it supports reset connection stored procedure.
2. This change is required for supporting RDS Proxy for Babelfish. RDS Proxy calls sys.sp_reset_connection before connection is reused by other clients.

Tasks:
BABEL-429

Signed off by: Sharath BP <[email protected]>

(cherry picked from commit 3997e2e)

contrib/babelfishpg_tds/src/backend/tds/tds_srv.c

@@ -197,6 +197,7 @@ pe_tds_init(void)
 	pltsql_plugin_handler_ptr->get_stat_values = &tds_stat_get_activity;
 	pltsql_plugin_handler_ptr->invalidate_stat_view = &invalidate_stat_table;
 	pltsql_plugin_handler_ptr->get_host_name = &get_tds_host_name;
+	pltsql_plugin_handler_ptr->set_reset_tds_connection_flag = &SetResetTDSConnectionFlag;
 
 	invalidate_stat_table_hook = invalidate_stat_table;
 	guc_newval_hook = TdsSetGucStatVariable;

contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c

@@ -73,6 +73,7 @@ typedef ResetConnectionData *ResetConnection;
 TdsRequestCtrlData *TdsRequestCtrl = NULL;
 
 ResetConnection resetCon = NULL;
+static bool resetTdsConnectionFlag = false;
 
 /* Local functions */
 static void ResetTDSConnection(void);
@@ -154,13 +155,28 @@ ResetTDSConnection(void)
 	TdsResetCache();
 	TdsResponseReset();
 	TdsResetBcpOffset();
-	SetConfigOption("default_transaction_isolation", isolationOld,
-					PGC_BACKEND, PGC_S_CLIENT);
+	/* Retore previous isolation level when not called by sys.sp_reset_connection */
+	if (!resetTdsConnectionFlag)
+	{
+		SetConfigOption("default_transaction_isolation", isolationOld,
+						PGC_BACKEND, PGC_S_CLIENT);
+	}
 
 	tvp_lookup_list = NIL;
 
-	/* send an environement change token */
-	TdsSendEnvChange(TDS_ENVID_RESETCON, NULL, NULL);
+	/* send an environement change token is its not called via sys.sp_reset_connection procedure */
+	if (!resetTdsConnectionFlag)
+	{
+		TdsSendEnvChange(TDS_ENVID_RESETCON, NULL, NULL);
+	}
+}
+
+/*
+ * SetResetTDSConnectionFlag - Sets reset tds connection flag
+ */
+void SetResetTDSConnectionFlag()
+{
+	resetTdsConnectionFlag = true;	
 }
 
 /*
@@ -674,6 +690,12 @@ TdsSocketBackend(void)
 						/* Ready to fetch the next request */
 						TdsRequestCtrl->phase = TDS_REQUEST_PHASE_FETCH;
 
+						if (resetTdsConnectionFlag)
+						{
+							ResetTDSConnection();
+							resetTdsConnectionFlag = false;
+						}
+
 						break;
 					}
 				case TDS_REQUEST_PHASE_ERROR:

contrib/babelfishpg_tds/src/include/tds_protocol.h

@@ -77,5 +77,6 @@ typedef struct
 } TdsRequestCtrlData;
 
 extern TdsRequestCtrlData *TdsRequestCtrl;
+extern void SetResetTDSConnectionFlag(void);
 
 #endif							/* TDS_PROTOCOL_H */

contrib/babelfishpg_tsql/sql/sys_procedures.sql

@@ -205,3 +205,7 @@ GRANT EXECUTE on PROCEDURE sys.sp_droprolemember(IN sys.SYSNAME, IN sys.SYSNAME)
 CREATE OR REPLACE PROCEDURE sys.sp_babelfish_volatility(IN "@function_name" sys.varchar DEFAULT NULL, IN "@volatility" sys.varchar DEFAULT NULL)
 AS 'babelfishpg_tsql', 'sp_babelfish_volatility' LANGUAGE C;
 GRANT EXECUTE on PROCEDURE sys.sp_babelfish_volatility(IN sys.varchar, IN sys.varchar) TO PUBLIC;
+
+CREATE OR REPLACE PROCEDURE sys.sp_reset_connection()
+AS 'babelfishpg_tsql', 'sp_reset_connection_internal' LANGUAGE C;
+GRANT EXECUTE ON PROCEDURE sys.sp_reset_connection() TO PUBLIC;

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--2.10.0--2.11.0.sql

@@ -17,3 +17,7 @@ CALL sys.analyze_babelfish_catalogs();
 
 -- Reset search_path to not affect any subsequent scripts
 SELECT set_config('search_path', trim(leading 'sys, ' from current_setting('search_path')), false);
+
+CREATE OR REPLACE PROCEDURE sys.sp_reset_connection()
+AS 'babelfishpg_tsql', 'sp_reset_connection_internal' LANGUAGE C;
+GRANT EXECUTE ON PROCEDURE sys.sp_reset_connection() TO PUBLIC;

contrib/babelfishpg_tsql/src/pltsql.h

@@ -1664,6 +1664,8 @@ typedef struct PLtsql_protocol_plugin
 
 	char	   *(*get_physical_schema_name) (char *db_name, const char *schema_name);
 
+	void		(*set_reset_tds_connection_flag) ();
+
 	/* Session level GUCs */
 	bool		quoted_identifier;
 	bool		arithabort;
@@ -1677,7 +1679,6 @@ typedef struct PLtsql_protocol_plugin
 	int			datefirst;
 	int			lock_timeout;
 	const char *language;
-
 } PLtsql_protocol_plugin;
 
 /*

contrib/babelfishpg_tsql/src/procedures.c

@@ -58,6 +58,7 @@ PG_FUNCTION_INFO_V1(sp_addrolemember);
 PG_FUNCTION_INFO_V1(sp_droprolemember);
 PG_FUNCTION_INFO_V1(sp_babelfish_volatility);
 PG_FUNCTION_INFO_V1(sp_rename_internal);
+PG_FUNCTION_INFO_V1(sp_reset_connection_internal);
 
 extern void delete_cached_batch(int handle);
 extern InlineCodeBlockArgs *create_args(int numargs);
@@ -2848,3 +2849,13 @@ gen_sp_rename_subcmds(const char *objname, const char *newname, const char *sche
 	return res;
 }
 
+Datum
+sp_reset_connection_internal(PG_FUNCTION_ARGS)
+{
+	if (*pltsql_protocol_plugin_ptr) 
+	{
+		(*pltsql_protocol_plugin_ptr)->set_reset_tds_connection_flag();
+	}
+
+	PG_RETURN_VOID();
+}

test/JDBC/expected/Test-sp_reset_connection.out

@@ -0,0 +1,109 @@
+-- 1. Test resets GUC variables
+SET lock_timeout 0;
+GO
+SELECT @@lock_timeout;
+GO
+~~START~~
+int
+0
+~~END~~
+
+EXEC sys.sp_reset_connection
+-- TODO: GUC is not resetting
+SELECT @@lock_timeout;
+GO
+~~START~~
+int
+0
+~~END~~
+
+
+-- 2. Test open transactions are aborted on reset
+DROP TABLE IF EXISTS sp_reset_connection_test_table;
+CREATE TABLE sp_reset_connection_test_table(id int);
+BEGIN TRANSACTION
+INSERT INTO sp_reset_connection_test_table VALUES(1)
+GO
+~~ROW COUNT: 1~~
+
+EXEC sys.sp_reset_connection
+GO
+COMMIT TRANSACTION
+GO
+~~ERROR (Code: 3902)~~
+
+~~ERROR (Message: COMMIT can only be used in transaction blocks)~~
+
+SELECT * FROM sp_reset_connection_test_table
+GO
+~~START~~
+int
+~~END~~
+
+
+-- 3. Test temp tables are deleted on reset
+CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+GO
+~~ROW COUNT: 3~~
+
+SELECT * from #babel_temp_table
+GO
+~~START~~
+int#!#int
+1#!#100
+2#!#200
+3#!#300
+~~END~~
+
+EXEC sys.sp_reset_connection
+GO
+SELECT * from #babel_temp_table
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: relation "#babel_temp_table" does not exist)~~
+
+
+-- 4. Test isolation level is reset
+SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+1
+~~END~~
+
+EXEC sys.sp_reset_connection
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+2
+~~END~~
+
+
+-- 5. Test sp_reset_connection called with sp_prepexec
+SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+1
+~~END~~
+
+DECLARE @handle int;
+EXEC SP_PREPARE @handle output, NULL, N'exec sys.sp_reset_connection'
+EXEC SP_EXECUTE @handle
+GO
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+2
+~~END~~
+

test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql

@@ -0,0 +1,56 @@
+-- 1. Test resets GUC variables
+SET lock_timeout 0;
+GO
+SELECT @@lock_timeout;
+GO
+EXEC sys.sp_reset_connection
+-- TODO: GUC is not resetting
+SELECT @@lock_timeout;
+GO
+
+-- 2. Test open transactions are aborted on reset
+DROP TABLE IF EXISTS sp_reset_connection_test_table;
+CREATE TABLE sp_reset_connection_test_table(id int);
+BEGIN TRANSACTION
+INSERT INTO sp_reset_connection_test_table VALUES(1)
+GO
+EXEC sys.sp_reset_connection
+GO
+COMMIT TRANSACTION
+GO
+SELECT * FROM sp_reset_connection_test_table
+GO
+
+-- 3. Test temp tables are deleted on reset
+CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+GO
+SELECT * from #babel_temp_table
+GO
+EXEC sys.sp_reset_connection
+GO
+SELECT * from #babel_temp_table
+GO
+
+-- 4. Test isolation level is reset
+SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+EXEC sys.sp_reset_connection
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+
+-- 5. Test sp_reset_connection called with sp_prepexec
+SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+DECLARE @handle int;
+EXEC SP_PREPARE @handle output, NULL, N'exec sys.sp_reset_connection'
+EXEC SP_EXECUTE @handle
+GO
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO

test/dotnet/ExpectedOutput/TestSpResetConnection.out

@@ -0,0 +1,36 @@
+#Q#DROP TABLE IF EXISTS sp_reset_connection_test_table;
+#Q#CREATE TABLE sp_reset_connection_test_table(id int);
+#Q#INSERT INTO sp_reset_connection_test_table VALUES(1)
+#Q#INSERT INTO sp_reset_connection_test_table VALUES(2)
+#Q#sys.sp_reset_connection
+#Q#SELECT * FROM sp_reset_connection_test_table
+#D#int
+1
+#Q#CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+#Q#INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+#Q#SELECT * from #babel_temp_table
+#D#int#!#int
+1#!#100
+2#!#200
+3#!#300
+#Q#sys.sp_reset_connection
+#Q#SELECT * from #babel_temp_table
+#E#relation "#babel_temp_table" does not exist
+#Q#select current_setting('transaction_isolation')
+#D#text
+repeatable read
+#Q#sys.sp_reset_connection
+#Q#select current_setting('transaction_isolation')
+#D#text
+read committed
+#Q#CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+#Q#INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+#Q#SELECT * from #babel_temp_table
+#D#int#!#int
+1#!#100
+2#!#200
+3#!#300
+#Q#sys.sp_reset_connection
+#Q#sys.sp_reset_connection
+#Q#SELECT * from #babel_temp_table
+#E#relation "#babel_temp_table" does not exist

test/dotnet/input/Storedproc/TestSpResetConnection.txt

@@ -0,0 +1,35 @@
+# Test (1): Test abort open transaction
+DROP TABLE IF EXISTS sp_reset_connection_test_table;
+CREATE TABLE sp_reset_connection_test_table(id int);
+INSERT INTO sp_reset_connection_test_table VALUES(1)
+txn#!#begin
+INSERT INTO sp_reset_connection_test_table VALUES(2)
+storedproc#!#prep#!#sys.sp_reset_connection#!#
+txn#!#commit
+SELECT * FROM sp_reset_connection_test_table
+# Test (1): End
+
+# Test (2): Test temp table deletion
+CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+SELECT * from #babel_temp_table
+storedproc#!#prep#!#sys.sp_reset_connection#!#
+SELECT * from #babel_temp_table
+# Test (2): End
+
+# Test (3): Test reset of isolation level
+txn#!#begin#!#isolation#!#ss
+select current_setting('transaction_isolation')
+storedproc#!#prep#!#sys.sp_reset_connection#!#
+select current_setting('transaction_isolation')
+txn#!#commit 
+# Test (3): End
+
+# Test (4): Test prepexec sp_reset_connection
+CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+SELECT * from #babel_temp_table
+prepst#!#sys.sp_reset_connection#!#
+prepst#!#exec#!#
+SELECT * from #babel_temp_table
+# Test (4): End

test/python/expected/sql_validation_framework/expected_create.out

@@ -192,6 +192,7 @@ Could not find upgrade tests for procedure sys.sp_databases
 Could not find upgrade tests for procedure sys.sp_describe_cursor
 Could not find upgrade tests for procedure sys.sp_oledb_ro_usrname
 Could not find upgrade tests for procedure sys.sp_prepare
+Could not find upgrade tests for procedure sys.sp_reset_connection
 Could not find upgrade tests for procedure sys.sp_unprepare
 Could not find upgrade tests for procedure sys.sp_updatestats
 Could not find upgrade tests for procedure xp_instance_regread