From 6d856fb0ad5b53729fabb2a247549b2487eb306e Mon Sep 17 00:00:00 2001
From: Sharath BP <sharatbp@amazon.com>
Date: Tue, 6 Aug 2024 08:15:33 +0000
Subject: [PATCH] Set the resetTDSConnectionFlag with inline handler in
 tdsprotocol instead of rondevouz variable

---
 contrib/babelfishpg_tds/src/backend/tds/tds_srv.c   |  3 +--
 .../babelfishpg_tds/src/backend/tds/tdsprotocol.c   | 12 +++++++++---
 contrib/babelfishpg_tds/src/include/tds_protocol.h  |  3 +--
 contrib/babelfishpg_tsql/src/pltsql.h               |  6 ++----
 contrib/babelfishpg_tsql/src/procedures.c           | 10 +---------
 test/JDBC/expected/Test-sp_reset_connection.out     | 13 +++----------
 6 files changed, 17 insertions(+), 30 deletions(-)

diff --git a/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c b/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c
index b4dde6feeb3..9bf5137aff2 100644
--- a/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c
+++ b/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c
@@ -202,8 +202,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->is_rpc_request = &isRPCRequest;
-	pltsql_plugin_handler_ptr->reset_tds_connection_flag = &resetTdsConnectionFlag;
+	pltsql_plugin_handler_ptr->set_reset_tds_connection_flag = &SetResetTDSConnectionFlag;
 
 	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
index e6ef9f6f995..5d1fd66ce68 100644
--- a/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c
+++ b/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c
@@ -71,10 +71,9 @@ typedef ResetConnectionData *ResetConnection;
  * Local structures
  */
 TdsRequestCtrlData *TdsRequestCtrl = NULL;
-bool resetTdsConnectionFlag;
-bool isRPCRequest;
 
 ResetConnection resetCon = NULL;
+static bool resetTdsConnectionFlag = false;
 
 /* Local functions */
 static void ResetTDSConnection(void);
@@ -164,6 +163,14 @@ ResetTDSConnection(void)
 	TdsSendEnvChange(TDS_ENVID_RESETCON, NULL, NULL);
 }
 
+/*
+ * SetResetTDSConnectionFlag - Sets reset tds connection flag
+ */
+void SetResetTDSConnectionFlag()
+{
+	resetTdsConnectionFlag = true;	
+}
+
 /*
  * GetTDSRequest - Fetch and parse a TDS packet and generate a TDS request that
  * can be processed later.
@@ -224,7 +231,6 @@ GetTDSRequest(bool *resetProtocol)
 				 */
 				HOLD_CANCEL_INTERRUPTS();
 				ret = TdsReadNextRequest(&message, &status, &messageType);
-				isRPCRequest = messageType == TDS_RPC;
 				RESUME_CANCEL_INTERRUPTS();
 
 				if (ret != 0)
diff --git a/contrib/babelfishpg_tds/src/include/tds_protocol.h b/contrib/babelfishpg_tds/src/include/tds_protocol.h
index 3994e1cc293..4018a65382e 100644
--- a/contrib/babelfishpg_tds/src/include/tds_protocol.h
+++ b/contrib/babelfishpg_tds/src/include/tds_protocol.h
@@ -77,7 +77,6 @@ typedef struct
 } TdsRequestCtrlData;
 
 extern TdsRequestCtrlData *TdsRequestCtrl;
-extern bool resetTdsConnectionFlag;
-extern bool isRPCRequest;
+extern void SetResetTDSConnectionFlag(void);
 
 #endif							/* TDS_PROTOCOL_H */
diff --git a/contrib/babelfishpg_tsql/src/pltsql.h b/contrib/babelfishpg_tsql/src/pltsql.h
index ac390d7580a..555c76b1cdd 100644
--- a/contrib/babelfishpg_tsql/src/pltsql.h
+++ b/contrib/babelfishpg_tsql/src/pltsql.h
@@ -1761,6 +1761,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;
@@ -1774,10 +1776,6 @@ 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
index d310f9c7b2c..d3a0781741d 100644
--- a/contrib/babelfishpg_tsql/src/procedures.c
+++ b/contrib/babelfishpg_tsql/src/procedures.c
@@ -4184,15 +4184,7 @@ sp_reset_connection_internal(PG_FUNCTION_ARGS)
 {
 	if (*pltsql_protocol_plugin_ptr) 
 	{
-		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'");
-		}
+		(*pltsql_protocol_plugin_ptr)->set_reset_tds_connection_flag();
 	}
 
 	PG_RETURN_VOID();
diff --git a/test/JDBC/expected/Test-sp_reset_connection.out b/test/JDBC/expected/Test-sp_reset_connection.out
index f9fbc2f1853..b0914bac684 100644
--- a/test/JDBC/expected/Test-sp_reset_connection.out
+++ b/test/JDBC/expected/Test-sp_reset_connection.out
@@ -16,17 +16,10 @@ int#!#int
 
 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~~
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: relation "#babel_temp_table" does not exist)~~