From cfb3122ebbd2f59f287c971096fb38820037c53d Mon Sep 17 00:00:00 2001 From: "Kushaalshroff1234@gmail.com" Date: Wed, 25 Sep 2024 11:43:12 +0000 Subject: [PATCH] Handle errors and add tests --- .../src/backend/tds/tdslogin.c | 145 +++++++++++------- .../src/backend/tds/tdsprotocol.c | 44 +++--- .../src/include/tds_protocol.h | 2 + contrib/babelfishpg_tsql/src/pl_exec-2.c | 17 +- contrib/babelfishpg_tsql/src/session.c | 2 - .../expected/Test-sp_reset_connection.out | 80 ++++++++++ .../Test-sp_reset_connection.sql | 56 ------- test/dotnet/ExpectedOutput/1_Setup.out | 6 + .../ExpectedOutput/1_Successful_reset.out | 5 + test/dotnet/src/ExecuteTests.cs | 2 + 10 files changed, 220 insertions(+), 139 deletions(-) delete mode 100644 test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql create mode 100644 test/dotnet/ExpectedOutput/1_Setup.out create mode 100644 test/dotnet/ExpectedOutput/1_Successful_reset.out diff --git a/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c b/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c index 6d053f0e2e6..089eb1b20c0 100644 --- a/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c +++ b/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c @@ -61,6 +61,7 @@ #include "src/include/tds_debug.h" #include "src/include/tds_int.h" +#include "src/include/tds_protocol.h" #include "src/include/tds_request.h" #include "src/include/tds_response.h" #include "src/include/guc.h" @@ -2049,6 +2050,16 @@ TdsProcessLogin(Port *port, bool loadedSsl) return rc; } +/* + * TdsSetDbContext: + * Used to Set the Database Context during login + * and during reset connection. + * Note: We should not optimize the scenario during + * reset connection to reset to the same database + * which might be in use since the USE db command + * will reset other configurations which might + * have changed. + */ void TdsSetDbContext() { @@ -2057,73 +2068,101 @@ TdsSetDbContext() char *user = NULL; MemoryContext oldContext = CurrentMemoryContext; - if (loginInfo->database != NULL && loginInfo->database[0] != '\0') + PG_TRY(); { - Oid db_id; + if (loginInfo->database != NULL && loginInfo->database[0] != '\0') + { + Oid db_id; + /* + * Before preparing the query, first check whether we got a valid + * database name and it exists. Otherwise, there'll be risk of + * SQL injection. + */ + StartTransactionCommand(); + db_id = pltsql_plugin_handler_ptr->pltsql_get_database_oid(loginInfo->database); + CommitTransactionCommand(); + MemoryContextSwitchTo(oldContext); + + if (!OidIsValid(db_id)) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_DATABASE), + errmsg("database \"%s\" does not exist", loginInfo->database))); + + /* + * Any delimitated/quoted db name identifier requested in login + * must be already handled before this point. + */ + useDbCommand = psprintf("USE [%s]", loginInfo->database); + dbname = pstrdup(loginInfo->database); + } + else + { + char *temp = NULL; + + StartTransactionCommand(); + temp = pltsql_plugin_handler_ptr->pltsql_get_login_default_db(loginInfo->username); + MemoryContextSwitchTo(oldContext); + + if (temp == NULL) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_DATABASE), + errmsg("could not find default database for user \"%s\"", loginInfo->username))); + + useDbCommand = psprintf("USE [%s]", temp); + dbname = pstrdup(temp); + CommitTransactionCommand(); + MemoryContextSwitchTo(oldContext); + } + + StartTransactionCommand(); /* - * Before preparing the query, first check whether we got a valid - * database name and it exists. Otherwise, there'll be risk of - * SQL injection. + * Check if user has privileges to access current database. */ - StartTransactionCommand(); - db_id = pltsql_plugin_handler_ptr->pltsql_get_database_oid(loginInfo->database); - CommitTransactionCommand(); - MemoryContextSwitchTo(oldContext); - - if (!OidIsValid(db_id)) + user = pltsql_plugin_handler_ptr->pltsql_get_user_for_database(dbname); + if (!user) ereport(ERROR, (errcode(ERRCODE_UNDEFINED_DATABASE), - errmsg("database \"%s\" does not exist", loginInfo->database))); + errmsg("Cannot open database \"%s\" requested by the login. The login failed", dbname))); /* - * Any delimitated/quoted db name identifier requested in login - * must be already handled before this point. + * loginInfo has a database name provided, so we execute a "USE + * []" through pltsql inline handler. */ - useDbCommand = psprintf("USE [%s]", loginInfo->database); - dbname = pstrdup(loginInfo->database); + ExecuteSQLBatch(useDbCommand); + CommitTransactionCommand(); } - else + PG_CATCH(); { - char *temp = NULL; - - StartTransactionCommand(); - temp = pltsql_plugin_handler_ptr->pltsql_get_login_default_db(loginInfo->username); - MemoryContextSwitchTo(oldContext); + /* + * If this is during reset phase and we encounter an error + * with mapped user or db not found then we should terminate + * the connection. + */ + if (resetTdsConnectionFlag) + { + /* Before terminating the connection, send the response to the client. */ + EmitErrorReport(); + FlushErrorState(); - if (temp == NULL) - ereport(ERROR, - (errcode(ERRCODE_UNDEFINED_DATABASE), - errmsg("could not find default database for user \"%s\"", loginInfo->username))); + /* + * Client driver terminates the connection with a + * dual error token and with error 596. Otherwise + * it sends the next requests before realising the + * session was terminated. + */ + TdsSendError(596, 1, ERROR, + "Cannot continue the execution because the session is in the kill state.", 1); - useDbCommand = psprintf("USE [%s]", temp); - dbname = pstrdup(temp); - CommitTransactionCommand(); - MemoryContextSwitchTo(oldContext); - } + TdsSendDone(TDS_TOKEN_DONE, TDS_DONE_ERROR, 0, 0); + TdsFlush(); - /* - * Check if user has privileges to access current database - */ - StartTransactionCommand(); - PG_TRY(); - { - user = pltsql_plugin_handler_ptr->pltsql_get_user_for_database(dbname); - if (!user) - ereport(ERROR, - (errcode(ERRCODE_UNDEFINED_DATABASE), - errmsg("Cannot open database \"%s\" requested by the login. The login failed", dbname))); - - /* - * loginInfo has a database name provided, so we execute a "USE - * []" through pgtsql inline handler - */ - ExecuteSQLBatch(useDbCommand); - CommitTransactionCommand(); - } - PG_CATCH(); - { - CommitTransactionCommand(); + /* Terminate the connection. */ + ereport(FATAL, + (errcode(ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION), + errmsg("Reset Connection Failed"))); + } + /* Else rethrow the error. */ PG_RE_THROW(); } PG_END_TRY(); diff --git a/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c b/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c index db9d0e48e53..27134fed262 100644 --- a/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c +++ b/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c @@ -73,10 +73,10 @@ typedef ResetConnectionData *ResetConnection; TdsRequestCtrlData *TdsRequestCtrl = NULL; ResetConnection resetCon = NULL; -static bool resetTdsConnectionFlag = false; +bool resetTdsConnectionFlag = false; /* Local functions */ -static void ResetTDSConnection(void); +static void ResetTDSConnection(); static TDSRequest GetTDSRequest(bool *resetProtocol); static void ProcessTDSRequest(TDSRequest request); static void enable_statement_timeout(void); @@ -119,7 +119,7 @@ TdsDiscardAll() * for RESETCON. */ static void -ResetTDSConnection(void) +ResetTDSConnection() { const char *isolationOld; @@ -161,7 +161,6 @@ ResetTDSConnection(void) SetConfigOption("default_transaction_isolation", isolationOld, PGC_BACKEND, PGC_S_CLIENT); } - TdsSetDbContext(); tvp_lookup_list = NIL; /* Send an environement change token is its not called via sys.sp_reset_connection procedure. */ @@ -295,7 +294,16 @@ GetTDSRequest(bool *resetProtocol) resetCon->messageType = messageType; resetCon->status = (status & ~TDS_PACKET_HEADER_STATUS_RESETCON); + /* + * Set resetTdsConnectionFlag to true so that we avoid + * sending any env change token for the USE DB command + * which will get executed. + */ + resetTdsConnectionFlag = true; + TdsSetDbContext(); + resetTdsConnectionFlag = false; ResetTDSConnection(); + TdsErrorContext->err_text = "Fetching TDS Request"; *resetProtocol = true; return NULL; @@ -364,23 +372,6 @@ GetTDSRequest(bool *resetProtocol) } PG_CATCH(); { - if (resetCon) - { - /* Before terminating the connection, send the response to the client */ - EmitErrorReport(); - FlushErrorState(); - - TdsSendError(596, 1, ERROR, - "Cannot continue the execution because the session is in the kill state.", 1); - - TdsSendDone(TDS_TOKEN_DONE, TDS_DONE_ERROR, 0, 0); - TdsFlush(); - - ereport(FATAL, - (errcode(ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION), - errmsg("Reset Connection Failed"))); - - } PG_RE_THROW(); } PG_END_TRY(); @@ -700,6 +691,17 @@ TdsSocketBackend(void) case TDS_REQUEST_PHASE_FLUSH: { TdsErrorContext->phase = "TDS_REQUEST_PHASE_FLUSH"; + + if (resetTdsConnectionFlag) + { + /* + * We must set the Db Context before resetting TDS state, + * becasue we need the existing TDS state to flush any errors + * along with the reset. + */ + TdsSetDbContext(); + } + /* Send the response now */ TdsFlush(); diff --git a/contrib/babelfishpg_tds/src/include/tds_protocol.h b/contrib/babelfishpg_tds/src/include/tds_protocol.h index 90eb9553b6a..7a80af2ceac 100644 --- a/contrib/babelfishpg_tds/src/include/tds_protocol.h +++ b/contrib/babelfishpg_tds/src/include/tds_protocol.h @@ -80,4 +80,6 @@ extern TdsRequestCtrlData *TdsRequestCtrl; extern void SetResetTDSConnectionFlag(void); extern bool GetResetTDSConnectionFlag(void); +extern bool resetTdsConnectionFlag; + #endif /* TDS_PROTOCOL_H */ diff --git a/contrib/babelfishpg_tsql/src/pl_exec-2.c b/contrib/babelfishpg_tsql/src/pl_exec-2.c index 75fcfbd5f11..503f14fd859 100644 --- a/contrib/babelfishpg_tsql/src/pl_exec-2.c +++ b/contrib/babelfishpg_tsql/src/pl_exec-2.c @@ -2916,15 +2916,18 @@ exec_stmt_usedb(PLtsql_execstate *estate, PLtsql_stmt_usedb *stmt) top_es_entry = top_es_entry->next; } + /* + * In case of reset-connection we do not need to send the environment change token. + */ if (!(*pltsql_protocol_plugin_ptr)->get_reset_tds_connection_flag()) { - snprintf(message, sizeof(message), "Changed database context to '%s'.", stmt->db_name); - /* send env change token to user */ - if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_env_change) - ((*pltsql_protocol_plugin_ptr)->send_env_change) (1, stmt->db_name, old_db_name); - /* send message to user */ - if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_info) - ((*pltsql_protocol_plugin_ptr)->send_info) (0, 1, 0, message, 0); + snprintf(message, sizeof(message), "Changed database context to '%s'.", stmt->db_name); + /* send env change token to user */ + if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_env_change) + ((*pltsql_protocol_plugin_ptr)->send_env_change) (1, stmt->db_name, old_db_name); + /* send message to user */ + if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_info) + ((*pltsql_protocol_plugin_ptr)->send_info) (0, 1, 0, message, 0); } return PLTSQL_RC_OK; } diff --git a/contrib/babelfishpg_tsql/src/session.c b/contrib/babelfishpg_tsql/src/session.c index 488ecf99b1e..aec26526333 100644 --- a/contrib/babelfishpg_tsql/src/session.c +++ b/contrib/babelfishpg_tsql/src/session.c @@ -207,8 +207,6 @@ void reset_session_properties(void) { reset_cached_batch(); - pltsql_explain_only = false; - pltsql_explain_analyze = false; } void diff --git a/test/JDBC/expected/Test-sp_reset_connection.out b/test/JDBC/expected/Test-sp_reset_connection.out index 0e3a80846c5..074da28281f 100644 --- a/test/JDBC/expected/Test-sp_reset_connection.out +++ b/test/JDBC/expected/Test-sp_reset_connection.out @@ -1,3 +1,4 @@ +-- tsql -- 1. Test resets GUC variables SET lock_timeout 0; GO @@ -107,3 +108,82 @@ smallint 2 ~~END~~ + +-- 6. Test Database Context being reset +-- Tests include negative cases where db is dropped or renamed +Create database reset_con_db1; +GO +Create database reset_con_db2; +GO +~~ERROR (Code: 33557097)~~ + +~~ERROR (Message: Database 'reset_con_db2' already exists. Choose a different database name.)~~ + + +-- tsql database=reset_con_db1 +select db_name(); +GO +~~START~~ +nvarchar +reset_con_db1 +~~END~~ + +exec sys.sp_reset_connection +GO +use master +GO +select db_name(); +GO +~~START~~ +nvarchar +master +~~END~~ + +exec sys.sp_reset_connection +GO +select db_name(); +GO +~~START~~ +nvarchar +reset_con_db1 +~~END~~ + +-- test db being dropped before resetting to same db +use master; +drop database reset_con_db1; +GO +exec sys.sp_reset_connection +GO +~~ERROR (Code: 911)~~ + +~~ERROR (Message: database "reset_con_db1" does not exist)~~ + +-- tsql database=reset_con_db2 +select db_name(); +GO +~~START~~ +nvarchar +reset_con_db2 +~~END~~ + +use master +GO +select db_name(); +GO +~~START~~ +nvarchar +master +~~END~~ + +ALTER DATABASE reset_con_db2 MODIFY NAME=reset_con_db3 +GO +exec sys.sp_reset_connection +GO +~~ERROR (Code: 911)~~ + +~~ERROR (Message: database "reset_con_db2" does not exist)~~ + + +-- tsql +DROP DATABASE reset_con_db3 +GO diff --git a/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql b/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql deleted file mode 100644 index d331eb9b498..00000000000 --- a/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql +++ /dev/null @@ -1,56 +0,0 @@ --- 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 diff --git a/test/dotnet/ExpectedOutput/1_Setup.out b/test/dotnet/ExpectedOutput/1_Setup.out new file mode 100644 index 00000000000..02405ed6ed9 --- /dev/null +++ b/test/dotnet/ExpectedOutput/1_Setup.out @@ -0,0 +1,6 @@ +#Q#create database db1; +#Q#use db1; +#Q#Select db_name() +#D#nvarchar +db1 + diff --git a/test/dotnet/ExpectedOutput/1_Successful_reset.out b/test/dotnet/ExpectedOutput/1_Successful_reset.out new file mode 100644 index 00000000000..98be25e6a6a --- /dev/null +++ b/test/dotnet/ExpectedOutput/1_Successful_reset.out @@ -0,0 +1,5 @@ +#Q#select db_name(); +#D#nvarchar +master +#Q#drop database db1; + diff --git a/test/dotnet/src/ExecuteTests.cs b/test/dotnet/src/ExecuteTests.cs index 40b212e85ab..3c326d36427 100644 --- a/test/dotnet/src/ExecuteTests.cs +++ b/test/dotnet/src/ExecuteTests.cs @@ -52,6 +52,8 @@ public void Test() } allFiles = tempList; } + allFiles = allFiles.OrderBy(file => file.DirectoryName) + .ThenBy(file => file.Name); Task[] tasksInParallel = new Task[allFiles.Count()]; bool [] result = new bool[allFiles.Count()]; int i = 0;