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]>

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

@@ -203,6 +203,7 @@ pe_tds_init(void)
 	pltsql_plugin_handler_ptr->set_context_info = &set_tds_context_info;
 	pltsql_plugin_handler_ptr->get_datum_from_byte_ptr = &TdsBytePtrToDatum;
 	pltsql_plugin_handler_ptr->get_datum_from_date_time_struct = &TdsDateTimeTypeToDatum;
+	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

@@ -329,3 +329,7 @@ GRANT EXECUTE ON PROCEDURE sys.sp_dropextendedproperty TO PUBLIC;
 CREATE OR REPLACE PROCEDURE sys.sp_enum_oledb_providers()
 AS 'babelfishpg_tsql', 'sp_enum_oledb_providers_internal' LANGUAGE C;
 GRANT EXECUTE on PROCEDURE sys.sp_enum_oledb_providers() 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--4.3.0--4.4.0.sql

@@ -1910,6 +1910,10 @@ CALL sys.babelfish_drop_deprecated_object('function', 'sys', 'sp_pkeys_internal_
 -- Please have this be one of the last statements executed in this upgrade script.
 DROP PROCEDURE sys.babelfish_drop_deprecated_object(varchar, varchar, varchar);
 
+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;
+
 -- After upgrade, always run analyze for all babelfish catalogs.
 CALL sys.analyze_babelfish_catalogs();
 

contrib/babelfishpg_tsql/src/pltsql.h

@@ -1797,6 +1797,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;
@@ -1810,7 +1812,6 @@ typedef struct PLtsql_protocol_plugin
 	int			datefirst;
 	int			lock_timeout;
 	const char *language;
-
 } PLtsql_protocol_plugin;
 
 /*

contrib/babelfishpg_tsql/src/procedures.c

@@ -78,6 +78,7 @@ PG_FUNCTION_INFO_V1(sp_babelfish_volatility);
 PG_FUNCTION_INFO_V1(sp_rename_internal);
 PG_FUNCTION_INFO_V1(sp_execute_postgresql);
 PG_FUNCTION_INFO_V1(sp_enum_oledb_providers_internal);
+PG_FUNCTION_INFO_V1(sp_reset_connection_internal);
 PG_FUNCTION_INFO_V1(sp_renamedb_internal);
 
 extern void delete_cached_batch(int handle);
@@ -4198,3 +4199,14 @@ sp_enum_oledb_providers_internal(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+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/JDBC/jdbc_schedule

@@ -383,4 +383,3 @@ ignore#!#string_agg-before-14_5-vu-cleanup
 ignore#!#string_agg_within-16_4-vu-prepare
 ignore#!#string_agg_within-16_4-vu-verify
 ignore#!#string_agg_within-16_4-vu-cleanup
-

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

@@ -190,6 +190,7 @@ Could not find upgrade tests for procedure sys.sp_cursor_list
 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 table sys.babelfish_configurations