diff --git a/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c b/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c index 63ce388c47a..b4dde6feeb3 100644 --- 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 index 038c244dda2..3bb8fedec63 100644 --- 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 index a661cc181f1..3994e1cc293 100644 --- 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 index 51b03784dfe..4bdbbf4bb58 100644 --- 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 index bbea96beec9..ac390d7580a 100644 --- 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 index 940fed0a438..8234f0d73bd 100644 --- 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 new file mode 100644 index 00000000000..f9fbc2f1853 --- /dev/null +++ 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 new file mode 100644 index 00000000000..0775618675b --- /dev/null +++ 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 index 89565031db8..759dff251b6 100644 --- 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 \ No newline at end of file +# 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 \ No newline at end of file