Add support for sys.sp_reset_connection stored procedure

babelfish-for-postgresql · Aug 2, 2024 · 7d70d95 · 7d70d95
1 parent 245d7ce
commit 7d70d95
Show file tree

Hide file tree

Showing 9 changed files with 87 additions and 2 deletions.
diff --git a/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c b/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c
@@ -202,6 +202,8 @@ 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->is_rpc_request = &isRPCRequest;
+	pltsql_plugin_handler_ptr->reset_tds_connection_flag = &resetTdsConnectionFlag;
 
 	invalidate_stat_table_hook = invalidate_stat_table;
 	guc_newval_hook = TdsSetGucStatVariable;

diff --git a/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c b/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c
@@ -71,6 +71,8 @@ typedef ResetConnectionData *ResetConnection;
  * Local structures
  */
 TdsRequestCtrlData *TdsRequestCtrl = NULL;
+bool resetTdsConnectionFlag;
+bool isRPCRequest;
 
 ResetConnection resetCon = NULL;
 
@@ -222,6 +224,7 @@ GetTDSRequest(bool *resetProtocol)
 				 */
 				HOLD_CANCEL_INTERRUPTS();
 				ret = TdsReadNextRequest(&message, &status, &messageType);
+				isRPCRequest = messageType == TDS_RPC;
 				RESUME_CANCEL_INTERRUPTS();
 
 				if (ret != 0)
@@ -261,7 +264,7 @@ GetTDSRequest(bool *resetProtocol)
 		 * memory context before exit so that we can process the request
 		 * later.
 		 */
-		if (status & TDS_PACKET_HEADER_STATUS_RESETCON)
+		if ((status & TDS_PACKET_HEADER_STATUS_RESETCON) || resetTdsConnectionFlag == true)
 		{
 			MemoryContextSwitchTo(TopMemoryContext);
 
@@ -276,6 +279,7 @@ GetTDSRequest(bool *resetProtocol)
 			ResetTDSConnection();
 			TdsErrorContext->err_text = "Fetching TDS Request";
 			*resetProtocol = true;
+			resetTdsConnectionFlag = false;
 			return NULL;
 		}
 
@@ -342,6 +346,7 @@ GetTDSRequest(bool *resetProtocol)
 	}
 	PG_CATCH();
 	{
+		resetTdsConnectionFlag = false;
 		PG_RE_THROW();
 	}
 	PG_END_TRY();

diff --git a/contrib/babelfishpg_tds/src/include/tds_protocol.h b/contrib/babelfishpg_tds/src/include/tds_protocol.h
@@ -77,5 +77,7 @@ typedef struct
 } TdsRequestCtrlData;
 
 extern TdsRequestCtrlData *TdsRequestCtrl;
+extern bool resetTdsConnectionFlag;
+extern bool isRPCRequest;
 
 #endif							/* TDS_PROTOCOL_H */
diff --git a/contrib/babelfishpg_tsql/sql/sys_procedures.sql b/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;
diff --git a/contrib/babelfishpg_tsql/src/pltsql.h b/contrib/babelfishpg_tsql/src/pltsql.h
@@ -1774,6 +1774,9 @@ typedef struct PLtsql_protocol_plugin
 	int			datefirst;
 	int			lock_timeout;
 	const char *language;
+
+	bool 		*reset_tds_connection_flag;
+	bool 		*is_rpc_request;
 
 } PLtsql_protocol_plugin;
 

diff --git a/contrib/babelfishpg_tsql/src/procedures.c b/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);
@@ -4177,3 +4178,19 @@ sp_enum_oledb_providers_internal(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+Datum
+sp_reset_connection_internal(PG_FUNCTION_ARGS)
+{
+	if (*(*pltsql_protocol_plugin_ptr)->is_rpc_request)
+	{
+		*(*pltsql_protocol_plugin_ptr)->reset_tds_connection_flag = true;
+	}
+	else
+	{
+		*(*pltsql_protocol_plugin_ptr)->reset_tds_connection_flag = false;
+		elog(ERROR, "Invalid object name 'sp_reset_connection'");
+	}
+
+	PG_RETURN_VOID();
+}
diff --git a/test/JDBC/expected/Test-sp_reset_connection.out b/test/JDBC/expected/Test-sp_reset_connection.out
@@ -0,0 +1,32 @@
+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
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Invalid object name 'sp_reset_connection')~~
+
+
+SELECT * from #babel_temp_table
+Go
+~~START~~
+int#!#int
+1#!#100
+2#!#200
+3#!#300
+~~END~~
+
diff --git a/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql b/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql
@@ -0,0 +1,12 @@
+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
diff --git a/test/dotnet/input/Storedproc/TestStoredProcedure.txt b/test/dotnet/input/Storedproc/TestStoredProcedure.txt
@@ -271,4 +271,12 @@ storedproc#!#prep#!#sp_test25#!#int|-|a|-|20|-|input#!#int|-|b|-|10|-|output#!#i
 storedproc#!#prep#!#sp_test25#!#int|-|a|-|20|-|input#!#int|-|b|-|10|-|inputoutput#!#int|-|c|-|10|-|input#!#int|-|d|-|10|-|output
 
 DROP PROCEDURE sp_test25
-# Test (24): End
+# Test (24): End
+
+# Test (25): Test sys.sp_reset_connection stored procedure
+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 (25): End