Handle errors and add tests

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
+		 * [<db_name>]" 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
-	 * [<db_name>]" 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();

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();
 

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 */

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;
 }

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

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